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 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 #[expect(
93 dead_code,
94 reason = "RAII guard: held solely so its Drop deregisters this view's pinned version. Never read."
95 )]
96 pin_guard: Option<crate::storage::adjacency_manager::PinGuard>,
97 fork_scope: Option<Arc<crate::fork::ForkScope>>,
105 backend: Arc<dyn StorageBackend>,
107 vid_labels_index: Arc<parking_lot::RwLock<crate::storage::vid_labels::VidLabelsIndex>>,
115 plugin_registry: Option<Arc<uni_plugin::PluginRegistry>>,
125}
126
127pub struct FlushInProgressGuard {
136 storage: Arc<StorageManager>,
137}
138
139impl FlushInProgressGuard {
140 pub fn new(storage: &Arc<StorageManager>) -> Self {
141 storage
142 .flush_in_progress
143 .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
144 Self {
145 storage: storage.clone(),
146 }
147 }
148}
149
150impl Drop for FlushInProgressGuard {
151 fn drop(&mut self) {
152 self.storage
154 .flush_in_progress
155 .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
156 }
157}
158
159fn is_lance_conflict(err: &anyhow::Error) -> bool {
165 let msg = err.to_string();
166 msg.contains("Incompatible transaction") || msg.contains("conflict")
167}
168
169async fn retry_on_lance_conflict<F, Fut>(mut op: F) -> anyhow::Result<()>
174where
175 F: FnMut() -> Fut,
176 Fut: std::future::Future<Output = anyhow::Result<()>>,
177{
178 for attempt in 0u32..10 {
179 match op().await {
180 Ok(()) => return Ok(()),
181 Err(e) => {
182 if !is_lance_conflict(&e) || attempt == 9 {
183 return Err(e);
184 }
185 let backoff_ms = 1u64 << attempt;
186 tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
187 }
188 }
189 }
190 unreachable!("retry loop exits via Ok or Err")
191}
192
193pub async fn merge_insert_batch_with_lance_conflict_retry(
202 backend: &dyn crate::backend::StorageBackend,
203 table_name: &str,
204 batch: arrow_array::RecordBatch,
205 on: &[&str],
206) -> anyhow::Result<()> {
207 retry_on_lance_conflict(|| async {
213 let exists = backend.table_exists(table_name).await?;
214 if !exists {
215 anyhow::bail!(
216 "merge_insert target table '{}' does not exist (partial writes \
217 require the row to already be present; CREATE goes through Append)",
218 table_name
219 );
220 }
221 backend
222 .merge_insert(table_name, on, vec![batch.clone()])
223 .await
224 })
225 .await
226}
227
228pub async fn write_batch_with_lance_conflict_retry(
237 backend: &dyn crate::backend::StorageBackend,
238 table_name: &str,
239 batch: arrow_array::RecordBatch,
240) -> anyhow::Result<()> {
241 use crate::backend::types::WriteMode;
242 retry_on_lance_conflict(|| async {
243 let exists = backend.table_exists(table_name).await?;
244 if exists {
245 backend
246 .write(table_name, vec![batch.clone()], WriteMode::Append)
247 .await
248 } else {
249 backend.create_table(table_name, vec![batch.clone()]).await
250 }
251 })
252 .await
253}
254
255struct CompactionGuard {
257 status: Arc<Mutex<CompactionStatus>>,
258}
259
260impl CompactionGuard {
261 fn new(status: Arc<Mutex<CompactionStatus>>) -> Option<Self> {
262 let mut s = acquire_mutex(&status, "compaction_status").ok()?;
263 if s.compaction_in_progress {
264 return None;
265 }
266 s.compaction_in_progress = true;
267 Some(Self {
268 status: status.clone(),
269 })
270 }
271}
272
273impl Drop for CompactionGuard {
274 fn drop(&mut self) {
275 match uni_common::sync::acquire_mutex(&self.status, "compaction_status") {
278 Ok(mut s) => {
279 s.compaction_in_progress = false;
280 s.last_compaction = Some(std::time::SystemTime::now());
281 }
282 Err(e) => {
283 log::error!(
287 "CompactionGuard drop failed to acquire poisoned lock: {}. \
288 Compaction status may be inconsistent. Issue #18/#150",
289 e
290 );
291 }
292 }
293 }
294}
295
296impl StorageManager {
297 pub async fn new_with_backend(
299 base_uri: &str,
300 store: Arc<dyn ObjectStore>,
301 backend: Arc<dyn StorageBackend>,
302 schema_manager: Arc<SchemaManager>,
303 config: UniConfig,
304 ) -> Result<Self> {
305 let resilient_store: Arc<dyn ObjectStore> = Arc::new(ResilientObjectStore::new(
306 store,
307 config.object_store.clone(),
308 ));
309
310 let snapshot_manager = Arc::new(SnapshotManager::new(resilient_store.clone()));
311
312 Self::recover_all_staging_tables(backend.as_ref(), &schema_manager).await?;
314
315 let mut sm = Self {
316 base_uri: base_uri.to_string(),
317 store: resilient_store,
318 schema_manager,
319 snapshot_manager,
320 adjacency_manager: Arc::new(AdjacencyManager::new(config.cache_size)),
321 config,
322 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
323 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
324 pinned_snapshot: None,
325 pinned_version_hwm: None,
326 pin_guard: None,
327 fork_scope: None,
328 backend,
329 vid_labels_index: Arc::new(parking_lot::RwLock::new(
330 crate::storage::vid_labels::VidLabelsIndex::new(),
331 )),
332 plugin_registry: None,
333 };
334
335 if let Err(e) = sm.rebuild_vid_labels_index().await {
339 warn!(
340 "Failed to rebuild VidLabelsIndex on startup: {}. Falling back to storage queries.",
341 e
342 );
343 }
344
345 Ok(sm)
346 }
347
348 #[cfg(feature = "lance-backend")]
350 pub async fn new(base_uri: &str, schema_manager: Arc<SchemaManager>) -> Result<Self> {
351 Self::new_with_config(base_uri, schema_manager, UniConfig::default()).await
352 }
353
354 #[cfg(feature = "lance-backend")]
356 pub async fn new_with_cache(
357 base_uri: &str,
358 schema_manager: Arc<SchemaManager>,
359 adjacency_cache_size: usize,
360 ) -> Result<Self> {
361 let config = UniConfig {
362 cache_size: adjacency_cache_size,
363 ..Default::default()
364 };
365 Self::new_with_config(base_uri, schema_manager, config).await
366 }
367
368 #[cfg(feature = "lance-backend")]
370 pub async fn new_with_config(
371 base_uri: &str,
372 schema_manager: Arc<SchemaManager>,
373 config: UniConfig,
374 ) -> Result<Self> {
375 let store = Self::build_store_from_uri(base_uri)?;
376 Self::new_with_store_and_config(base_uri, store, schema_manager, config).await
377 }
378
379 #[cfg(feature = "lance-backend")]
381 pub async fn new_with_store_and_config(
382 base_uri: &str,
383 store: Arc<dyn ObjectStore>,
384 schema_manager: Arc<SchemaManager>,
385 config: UniConfig,
386 ) -> Result<Self> {
387 Self::new_with_store_and_storage_options(base_uri, store, schema_manager, config, None)
388 .await
389 }
390
391 #[cfg(feature = "lance-backend")]
393 pub async fn new_with_store_and_storage_options(
394 base_uri: &str,
395 store: Arc<dyn ObjectStore>,
396 schema_manager: Arc<SchemaManager>,
397 config: UniConfig,
398 lancedb_storage_options: Option<HashMap<String, String>>,
399 ) -> Result<Self> {
400 let backend = Arc::new(LanceDbBackend::connect(base_uri, lancedb_storage_options).await?);
401 Self::new_with_backend(base_uri, store, backend, schema_manager, config).await
402 }
403
404 async fn recover_all_staging_tables(
409 backend: &dyn StorageBackend,
410 schema_manager: &SchemaManager,
411 ) -> Result<()> {
412 let schema = schema_manager.schema();
413
414 backend
416 .recover_staging(table_names::main_vertex_table_name())
417 .await?;
418 backend
419 .recover_staging(table_names::main_edge_table_name())
420 .await?;
421
422 for label in schema.labels.keys() {
424 let name = table_names::vertex_table_name(label);
425 backend.recover_staging(&name).await?;
426 }
427
428 for edge_type in schema.edge_types.keys() {
430 for direction in &["fwd", "bwd"] {
431 let delta_name = table_names::delta_table_name(edge_type, direction);
433 backend.recover_staging(&delta_name).await?;
434
435 for _label in schema.labels.keys() {
437 let adj_name = table_names::adjacency_table_name(edge_type, direction);
438 backend.recover_staging(&adj_name).await?;
439 }
440 }
441 }
442
443 Ok(())
444 }
445
446 #[cfg(feature = "lance-backend")]
447 fn build_store_from_uri(base_uri: &str) -> Result<Arc<dyn ObjectStore>> {
448 if base_uri.contains("://") {
449 let parsed = url::Url::parse(base_uri).map_err(|e| anyhow!("Invalid base URI: {e}"))?;
450 let (store, _path) = object_store::parse_url(&parsed)
451 .map_err(|e| anyhow!("Failed to parse object store URL: {e}"))?;
452 Ok(Arc::from(store))
453 } else {
454 std::fs::create_dir_all(base_uri)?;
456 Ok(Arc::new(LocalFileSystem::new_with_prefix(base_uri)?))
457 }
458 }
459
460 pub fn local_fs_root(&self) -> Option<std::path::PathBuf> {
466 if self.base_uri.contains("://") {
467 None
468 } else {
469 Some(std::path::PathBuf::from(&self.base_uri))
470 }
471 }
472
473 pub fn pinned(&self, snapshot: SnapshotManifest) -> Self {
474 Self {
482 base_uri: self.base_uri.clone(),
483 store: self.store.clone(),
484 schema_manager: self.schema_manager.clone(),
485 snapshot_manager: self.snapshot_manager.clone(),
486 adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
490 config: self.config.clone(),
491 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
492 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
493 pinned_snapshot: Some(snapshot),
494 pinned_version_hwm: None,
495 pin_guard: None,
496 fork_scope: self.fork_scope.clone(),
497 backend: self.backend.clone(),
498 vid_labels_index: Arc::new(parking_lot::RwLock::new(
504 self.vid_labels_index.read().clone(),
505 )),
506 plugin_registry: self.plugin_registry.clone(),
507 }
508 }
509
510 pub fn pinned_at_version(&self, hwm: u64) -> Self {
529 let pin_guard = Some(self.adjacency_manager.pinned_versions().pin(hwm));
532 Self {
533 base_uri: self.base_uri.clone(),
534 store: self.store.clone(),
535 schema_manager: self.schema_manager.clone(),
536 snapshot_manager: self.snapshot_manager.clone(),
537 adjacency_manager: self.adjacency_manager.clone(),
538 config: self.config.clone(),
539 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
540 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
541 pinned_snapshot: None,
542 pinned_version_hwm: Some(hwm),
543 pin_guard,
544 fork_scope: self.fork_scope.clone(),
545 backend: self.backend.clone(),
546 vid_labels_index: Arc::new(parking_lot::RwLock::new(
552 self.vid_labels_index.read().clone(),
553 )),
554 plugin_registry: self.plugin_registry.clone(),
555 }
556 }
557
558 pub fn at_fork(&self, scope: Arc<crate::fork::ForkScope>) -> Self {
573 self.at_fork_with_schema(scope, self.schema_manager.clone())
574 }
575
576 pub fn at_fork_with_schema(
584 &self,
585 scope: Arc<crate::fork::ForkScope>,
586 merged_schema: Arc<SchemaManager>,
587 ) -> Self {
588 debug_assert!(
589 self.pinned_snapshot.is_none(),
590 "forking a pinned StorageManager is unsupported in Phase 1"
591 );
592 let branched_backend: Arc<dyn StorageBackend> = Arc::new(
593 crate::backend::branched::BranchedBackend::new(self.backend.clone(), scope.clone()),
594 );
595 let snapshot_manager = Arc::new(SnapshotManager::new_for_fork(
600 self.store.clone(),
601 scope.fork_id(),
602 ));
603 Self {
604 base_uri: self.base_uri.clone(),
605 store: self.store.clone(),
606 schema_manager: merged_schema,
607 snapshot_manager,
608 adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
609 config: self.config.clone(),
610 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
611 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
612 pinned_snapshot: None,
613 pinned_version_hwm: None,
614 pin_guard: None,
615 fork_scope: Some(scope),
616 backend: branched_backend,
617 vid_labels_index: Arc::new(parking_lot::RwLock::new(
623 self.vid_labels_index.read().clone(),
624 )),
625 plugin_registry: self.plugin_registry.clone(),
626 }
627 }
628
629 pub fn plugin_registry(&self) -> Option<&Arc<uni_plugin::PluginRegistry>> {
634 self.plugin_registry.as_ref()
635 }
636
637 pub fn set_plugin_registry(&mut self, registry: Arc<uni_plugin::PluginRegistry>) {
643 self.plugin_registry = Some(registry);
644 }
645
646 pub fn fork_scope(&self) -> Option<&Arc<crate::fork::ForkScope>> {
648 self.fork_scope.as_ref()
649 }
650
651 #[must_use]
664 pub fn has_fork_index(
665 &self,
666 label: &str,
667 column: &str,
668 kind: crate::fork::ForkLocalIndexKind,
669 ) -> bool {
670 self.fork_scope
671 .as_ref()
672 .is_some_and(|s| s.has_fork_local_index(label, column, kind))
673 }
674
675 pub fn base_uri(&self) -> &str {
678 &self.base_uri
679 }
680
681 pub fn get_edge_version_by_id(&self, edge_type_id: u32) -> Option<u64> {
682 let schema = self.schema_manager.schema();
683 let name = schema.edge_type_name_by_id(edge_type_id)?;
684 self.pinned_snapshot
685 .as_ref()
686 .and_then(|s| s.edges.get(name).map(|es| es.lance_version))
687 .filter(|v| *v != 0)
692 }
693
694 pub fn version_high_water_mark(&self) -> Option<u64> {
702 self.pinned_snapshot
703 .as_ref()
704 .map(|s| s.version_high_water_mark)
705 .or(self.pinned_version_hwm)
706 }
707
708 pub fn snapshot_version_hwm(&self) -> Option<u64> {
720 self.pinned_snapshot
721 .as_ref()
722 .map(|s| s.version_high_water_mark)
723 }
724
725 pub fn apply_version_filter(&self, base_filter: FilterExpr) -> FilterExpr {
730 match self.version_high_water_mark() {
731 Some(hwm) => FilterExpr::all([base_filter, FilterExpr::version_at_most(hwm)]),
732 None => base_filter,
733 }
734 }
735
736 fn build_active_filter(user_filter: Option<&str>) -> FilterExpr {
742 match user_filter {
743 Some(expr) => {
744 FilterExpr::all([FilterExpr::Raw(expr.to_string()), FilterExpr::not_deleted()])
745 }
746 None => FilterExpr::not_deleted(),
747 }
748 }
749
750 pub fn store(&self) -> Arc<dyn ObjectStore> {
751 self.store.clone()
752 }
753
754 pub fn compaction_status(
760 &self,
761 ) -> Result<CompactionStatus, uni_common::sync::LockPoisonedError> {
762 let guard = uni_common::sync::acquire_mutex(&self.compaction_status, "compaction_status")?;
763 Ok(guard.clone())
764 }
765
766 pub async fn compact(&self) -> Result<CompactionStats> {
767 let start = std::time::Instant::now();
769 let schema = self.schema_manager.schema();
770 let mut files_compacted = 0;
771
772 for label in schema.labels.keys() {
773 let name = table_names::vertex_table_name(label);
774 if self.backend.table_exists(&name).await? {
775 self.backend.optimize_table(&name).await?;
776 files_compacted += 1;
777 self.backend.invalidate_cache(&name);
778 }
779 }
780
781 Ok(CompactionStats {
782 files_compacted,
783 bytes_before: 0,
784 bytes_after: 0,
785 duration: start.elapsed(),
786 crdt_merges: 0,
787 })
788 }
789
790 pub async fn compact_label(&self, label: &str) -> Result<CompactionStats> {
791 let _guard = CompactionGuard::new(self.compaction_status.clone())
792 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
793
794 let start = std::time::Instant::now();
795 let name = table_names::vertex_table_name(label);
796
797 if self.backend.table_exists(&name).await? {
798 self.backend.optimize_table(&name).await?;
799 self.backend.invalidate_cache(&name);
800 }
801
802 Ok(CompactionStats {
803 files_compacted: 1,
804 bytes_before: 0,
805 bytes_after: 0,
806 duration: start.elapsed(),
807 crdt_merges: 0,
808 })
809 }
810
811 pub async fn compact_edge_type(&self, edge_type: &str) -> Result<CompactionStats> {
812 let _guard = CompactionGuard::new(self.compaction_status.clone())
813 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
814
815 let start = std::time::Instant::now();
816 let mut files_compacted = 0;
817
818 for dir in ["fwd", "bwd"] {
819 let name = table_names::delta_table_name(edge_type, dir);
820 if self.backend.table_exists(&name).await? {
821 self.backend.optimize_table(&name).await?;
822 files_compacted += 1;
823 }
824 }
825
826 Ok(CompactionStats {
827 files_compacted,
828 bytes_before: 0,
829 bytes_after: 0,
830 duration: start.elapsed(),
831 crdt_merges: 0,
832 })
833 }
834
835 pub async fn wait_for_compaction(&self) -> Result<()> {
836 loop {
837 let in_progress = {
838 acquire_mutex(&self.compaction_status, "compaction_status")?.compaction_in_progress
839 };
840 if !in_progress {
841 return Ok(());
842 }
843 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
844 }
845 }
846
847 pub fn start_background_compaction(
848 self: Arc<Self>,
849 mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
850 ) -> tokio::task::JoinHandle<()> {
851 if !self.config.compaction.enabled {
852 return tokio::spawn(async {});
853 }
854
855 tokio::spawn(async move {
856 let start = tokio::time::Instant::now() + self.config.compaction.check_interval;
862 let mut interval =
863 tokio::time::interval_at(start, self.config.compaction.check_interval);
864
865 loop {
866 tokio::select! {
867 _ = interval.tick() => {
868 if let Err(e) = self.update_compaction_status().await {
869 log::error!("Failed to update compaction status: {}", e);
870 continue;
871 }
872
873 if let Some(task) = self.pick_compaction_task() {
874 log::info!("Triggering background compaction: {:?}", task);
875 if let Err(e) = Self::execute_compaction(Arc::clone(&self), task).await {
876 log::error!("Compaction failed: {}", e);
877 }
878 }
879 }
880 _ = shutdown_rx.recv() => {
881 log::info!("Background compaction shutting down");
882 let _ = self.wait_for_compaction().await;
883 break;
884 }
885 }
886 }
887 })
888 }
889
890 async fn update_compaction_status(&self) -> Result<()> {
891 let schema = self.schema_manager.schema();
892 let backend = self.backend.as_ref();
893 let mut total_rows: usize = 0;
894 let mut oldest_ts: Option<i64> = None;
895
896 for name in schema.edge_types.keys() {
897 for dir in ["fwd", "bwd"] {
898 let tbl_name = table_names::delta_table_name(name, dir);
899 if !backend.table_exists(&tbl_name).await? {
900 continue;
901 }
902 let row_count = backend.count_rows(&tbl_name, None).await.unwrap_or(0);
903 if row_count == 0 {
904 continue;
905 }
906 total_rows += row_count;
907
908 let request =
910 ScanRequest::all(&tbl_name).with_columns(vec!["_created_at".to_string()]);
911 let Ok(batches) = backend.scan(request).await else {
912 continue;
913 };
914 for batch in batches {
915 let Some(col) = batch
916 .column_by_name("_created_at")
917 .and_then(|c| c.as_any().downcast_ref::<TimestampNanosecondArray>())
918 else {
919 continue;
920 };
921 for i in 0..col.len() {
922 if !col.is_null(i) {
923 let ts = col.value(i);
924 oldest_ts = Some(oldest_ts.map_or(ts, |prev| prev.min(ts)));
925 }
926 }
927 }
928 }
929 }
930
931 let oldest_l1_age = oldest_ts
932 .and_then(|ts| {
933 let created = UNIX_EPOCH + Duration::from_nanos(ts as u64);
934 SystemTime::now().duration_since(created).ok()
935 })
936 .unwrap_or(Duration::ZERO);
937
938 let mut status = acquire_mutex(&self.compaction_status, "compaction_status")?;
939 status.l1_size_bytes = (total_rows * ENTRY_SIZE_ESTIMATE) as u64;
942 status.oldest_l1_age = oldest_l1_age;
943 Ok(())
944 }
945
946 fn pick_compaction_task(&self) -> Option<CompactionTask> {
947 let status = acquire_mutex(&self.compaction_status, "compaction_status").ok()?;
948
949 if status.l1_runs >= self.config.compaction.max_l1_runs {
950 return Some(CompactionTask::ByRunCount);
951 }
952 if status.l1_size_bytes >= self.config.compaction.max_l1_size_bytes {
953 return Some(CompactionTask::BySize);
954 }
955 if status.oldest_l1_age >= self.config.compaction.max_l1_age
956 && status.oldest_l1_age > Duration::ZERO
957 {
958 return Some(CompactionTask::ByAge);
959 }
960
961 None
962 }
963
964 async fn try_optimize_table(backend: &dyn StorageBackend, table_name: &str) -> bool {
966 if let Err(e) = backend.optimize_table(table_name).await {
967 log::warn!("Failed to optimize table {}: {}", table_name, e);
968 return false;
969 }
970 true
971 }
972
973 pub fn trigger_async_compaction(self: &Arc<Self>) {
976 let this = Arc::clone(self);
977 tokio::spawn(async move {
978 if let Err(e) = Self::execute_compaction(this, CompactionTask::ByRunCount).await {
979 log::debug!("Post-flush compaction skipped: {}", e);
981 }
982 });
983 }
984
985 pub(crate) async fn execute_compaction(
986 this: Arc<Self>,
987 _task: CompactionTask,
988 ) -> Result<CompactionStats> {
989 let start = std::time::Instant::now();
990 let _guard = CompactionGuard::new(this.compaction_status.clone())
991 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
992
993 let schema = this.schema_manager.schema();
994 let mut files_compacted = 0;
995
996 let compactor = Compactor::new(Arc::clone(&this));
999 let compaction_results = compactor.compact_all().await.unwrap_or_else(|e| {
1000 log::error!(
1001 "Semantic compaction failed (continuing with backend optimize): {}",
1002 e
1003 );
1004 Vec::new()
1005 });
1006
1007 let am = this.adjacency_manager();
1009 for info in &compaction_results {
1010 let direction = match info.direction.as_str() {
1011 "fwd" => Direction::Outgoing,
1012 "bwd" => Direction::Incoming,
1013 _ => continue,
1014 };
1015 if let Some(etid) = schema.edge_type_id_unified_case_insensitive(&info.edge_type)
1016 && let Err(e) = am.warm(&this, etid, direction, None).await
1017 {
1018 log::warn!(
1019 "Failed to re-warm adjacency for {}/{}: {}",
1020 info.edge_type,
1021 info.direction,
1022 e
1023 );
1024 }
1025 }
1026
1027 let backend = this.backend.as_ref();
1029
1030 for name in schema.edge_types.keys() {
1032 for dir in ["fwd", "bwd"] {
1033 let delta = table_names::delta_table_name(name, dir);
1034 if Self::try_optimize_table(backend, &delta).await {
1035 files_compacted += 1;
1036 }
1037 let adj = table_names::adjacency_table_name(name, dir);
1038 if Self::try_optimize_table(backend, &adj).await {
1039 files_compacted += 1;
1040 }
1041 }
1042 }
1043
1044 for label in schema.labels.keys() {
1046 let tbl = table_names::vertex_table_name(label);
1047 if Self::try_optimize_table(backend, &tbl).await {
1048 files_compacted += 1;
1049 backend.invalidate_cache(&tbl);
1050 }
1051 }
1052
1053 for tbl in [
1055 table_names::main_vertex_table_name(),
1056 table_names::main_edge_table_name(),
1057 ] {
1058 if Self::try_optimize_table(backend, tbl).await {
1059 files_compacted += 1;
1060 }
1061 }
1062
1063 {
1064 let mut status = acquire_mutex(&this.compaction_status, "compaction_status")?;
1065 status.total_compactions += 1;
1066 status.l1_runs = 0; }
1068
1069 Ok(CompactionStats {
1070 files_compacted,
1071 bytes_before: 0,
1072 bytes_after: 0,
1073 duration: start.elapsed(),
1074 crdt_merges: 0,
1075 })
1076 }
1077
1078 pub fn invalidate_table_cache(&self, label: &str) {
1082 let name = table_names::vertex_table_name(label);
1083 self.backend.invalidate_cache(&name);
1084 }
1085
1086 pub fn base_path(&self) -> &str {
1087 &self.base_uri
1088 }
1089
1090 pub fn schema_manager(&self) -> &SchemaManager {
1091 &self.schema_manager
1092 }
1093
1094 pub fn schema_manager_arc(&self) -> Arc<SchemaManager> {
1095 self.schema_manager.clone()
1096 }
1097
1098 #[must_use]
1106 pub fn schema_manager_arc_ref(&self) -> &Arc<SchemaManager> {
1107 &self.schema_manager
1108 }
1109
1110 pub fn adjacency_manager(&self) -> Arc<AdjacencyManager> {
1112 Arc::clone(&self.adjacency_manager)
1113 }
1114
1115 pub async fn warm_adjacency(
1120 &self,
1121 edge_type_id: u32,
1122 direction: crate::storage::direction::Direction,
1123 version: Option<u64>,
1124 ) -> anyhow::Result<()> {
1125 self.adjacency_manager
1126 .warm(self, edge_type_id, direction, version)
1127 .await
1128 }
1129
1130 pub async fn warm_adjacency_coalesced(
1135 &self,
1136 edge_type_id: u32,
1137 direction: crate::storage::direction::Direction,
1138 version: Option<u64>,
1139 ) -> anyhow::Result<()> {
1140 self.adjacency_manager
1141 .warm_coalesced(self, edge_type_id, direction, version)
1142 .await
1143 }
1144
1145 pub fn has_adjacency_csr(
1147 &self,
1148 edge_type_id: u32,
1149 direction: crate::storage::direction::Direction,
1150 ) -> bool {
1151 self.adjacency_manager.has_csr(edge_type_id, direction)
1152 }
1153
1154 pub fn get_neighbors_at_version(
1156 &self,
1157 vid: uni_common::core::id::Vid,
1158 edge_type: u32,
1159 direction: crate::storage::direction::Direction,
1160 version: u64,
1161 ) -> Vec<(uni_common::core::id::Vid, uni_common::core::id::Eid)> {
1162 self.adjacency_manager
1163 .get_neighbors_at_version(vid, edge_type, direction, version)
1164 }
1165
1166 pub fn backend(&self) -> &dyn StorageBackend {
1168 self.backend.as_ref()
1169 }
1170
1171 pub fn backend_arc(&self) -> Arc<dyn StorageBackend> {
1173 self.backend.clone()
1174 }
1175
1176 async fn rebuild_vid_labels_index(&mut self) -> Result<()> {
1181 use crate::storage::vid_labels::VidLabelsIndex;
1182
1183 let backend = self.backend.as_ref();
1184 let vtable = table_names::main_vertex_table_name();
1185
1186 if !backend.table_exists(vtable).await? {
1188 self.vid_labels_index = Arc::new(parking_lot::RwLock::new(VidLabelsIndex::new()));
1189 return Ok(());
1190 }
1191
1192 let request = ScanRequest::all(vtable)
1194 .with_filter(FilterExpr::not_deleted())
1195 .with_limit(100_000);
1196 let batches = backend
1197 .scan(request)
1198 .await
1199 .map_err(|e| anyhow!("Failed to query main vertex table: {}", e))?;
1200
1201 let mut index = VidLabelsIndex::new();
1202 for batch in batches {
1203 let vid_col = batch
1204 .column_by_name("_vid")
1205 .ok_or_else(|| anyhow!("Missing _vid column"))?
1206 .as_any()
1207 .downcast_ref::<UInt64Array>()
1208 .ok_or_else(|| anyhow!("Invalid _vid column type"))?;
1209
1210 let labels_col = batch
1211 .column_by_name("labels")
1212 .ok_or_else(|| anyhow!("Missing labels column"))?
1213 .as_any()
1214 .downcast_ref::<arrow_array::ListArray>()
1215 .ok_or_else(|| anyhow!("Invalid labels column type"))?;
1216
1217 for row_idx in 0..batch.num_rows() {
1218 let vid = Vid::from(vid_col.value(row_idx));
1219 let labels_array = labels_col.value(row_idx);
1220 let labels_str_array = labels_array
1221 .as_any()
1222 .downcast_ref::<arrow_array::StringArray>()
1223 .ok_or_else(|| anyhow!("Invalid labels array element type"))?;
1224
1225 let labels: Vec<String> = (0..labels_str_array.len())
1226 .map(|i| labels_str_array.value(i).to_string())
1227 .collect();
1228
1229 index.insert(vid, labels);
1230 }
1231 }
1232
1233 self.vid_labels_index = Arc::new(parking_lot::RwLock::new(index));
1234 Ok(())
1235 }
1236
1237 pub fn get_labels_from_index(&self, vid: Vid) -> Option<Vec<String>> {
1242 let index = self.vid_labels_index.read();
1243 index.get_labels(vid).map(|labels| labels.to_vec())
1244 }
1245
1246 pub fn physical_vertex_label_names(&self) -> Vec<String> {
1253 self.vid_labels_index
1254 .read()
1255 .labels()
1256 .map(str::to_owned)
1257 .collect()
1258 }
1259
1260 pub fn update_vid_labels_index(&self, vid: Vid, labels: Vec<String>) {
1262 let mut index = self.vid_labels_index.write();
1263 index.insert(vid, labels);
1264 }
1265
1266 pub fn remove_from_vid_labels_index(&self, vid: Vid) {
1268 let mut index = self.vid_labels_index.write();
1269 index.remove_vid(vid);
1270 }
1271
1272 pub async fn load_subgraph_cached(
1273 &self,
1274 start_vids: &[Vid],
1275 edge_types: &[u32],
1276 max_hops: usize,
1277 direction: GraphDirection,
1278 _l0: Option<Arc<RwLock<L0Buffer>>>,
1279 ) -> Result<WorkingGraph> {
1280 let mut graph = WorkingGraph::new();
1281
1282 let dir = match direction {
1283 GraphDirection::Outgoing => crate::storage::direction::Direction::Outgoing,
1284 GraphDirection::Incoming => crate::storage::direction::Direction::Incoming,
1285 };
1286
1287 let neighbor_is_dst = matches!(direction, GraphDirection::Outgoing);
1288
1289 let mut frontier: Vec<Vid> = start_vids.to_vec();
1291 let mut visited: HashSet<Vid> = HashSet::new();
1292
1293 for &vid in start_vids {
1295 graph.add_vertex(vid);
1296 }
1297
1298 for _hop in 0..max_hops {
1299 let mut next_frontier = HashSet::new();
1300
1301 for &vid in &frontier {
1302 if visited.contains(&vid) {
1303 continue;
1304 }
1305 visited.insert(vid);
1306 graph.add_vertex(vid);
1307
1308 for &etype_id in edge_types {
1309 let edge_ver = self.snapshot_version_hwm();
1313 self.adjacency_manager
1314 .warm_coalesced(self, etype_id, dir, edge_ver)
1315 .await?;
1316
1317 let edges = self.adjacency_manager.get_neighbors(vid, etype_id, dir);
1319
1320 for (neighbor_vid, eid) in edges {
1321 graph.add_vertex(neighbor_vid);
1322 if !visited.contains(&neighbor_vid) {
1323 next_frontier.insert(neighbor_vid);
1324 }
1325
1326 if neighbor_is_dst {
1327 graph.add_edge(vid, neighbor_vid, eid, etype_id);
1328 } else {
1329 graph.add_edge(neighbor_vid, vid, eid, etype_id);
1330 }
1331 }
1332 }
1333 }
1334 frontier = next_frontier.into_iter().collect();
1335
1336 if frontier.is_empty() {
1338 break;
1339 }
1340 }
1341
1342 Ok(graph)
1343 }
1344
1345 pub fn snapshot_manager(&self) -> &SnapshotManager {
1346 &self.snapshot_manager
1347 }
1348
1349 pub fn index_manager(&self) -> IndexManager {
1350 IndexManager::new(&self.base_uri, self.schema_manager.clone())
1351 .with_backend(self.backend_arc())
1352 }
1353
1354 pub async fn scan_vertex_table(
1363 &self,
1364 label: &str,
1365 columns: &[&str],
1366 additional_filter: Option<&FilterExpr>,
1367 ) -> Result<Option<arrow_array::RecordBatch>> {
1368 let backend = self.backend();
1369 let table_name = table_names::vertex_table_name(label);
1370
1371 if !backend.table_exists(&table_name).await? {
1372 return Ok(None);
1373 }
1374
1375 let actual_columns =
1377 if let Some(table_schema) = backend.get_table_schema(&table_name).await? {
1378 let table_field_names: HashSet<&str> = table_schema
1379 .fields()
1380 .iter()
1381 .map(|f| f.name().as_str())
1382 .collect();
1383 columns
1384 .iter()
1385 .copied()
1386 .filter(|c| table_field_names.contains(c))
1387 .map(|s| s.to_string())
1388 .collect::<Vec<_>>()
1389 } else {
1390 return Ok(None);
1391 };
1392
1393 let filter = match (self.version_high_water_mark(), additional_filter) {
1395 (Some(hwm), Some(f)) => Some(FilterExpr::all([
1396 FilterExpr::version_at_most(hwm),
1397 f.clone(),
1398 ])),
1399 (Some(hwm), None) => Some(FilterExpr::version_at_most(hwm)),
1400 (None, Some(f)) => Some(f.clone()),
1401 (None, None) => None,
1402 };
1403
1404 let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1405 if let Some(f) = filter {
1406 request = request.with_filter(f);
1407 }
1408
1409 let batches = backend.scan(request).await?;
1416 if batches.is_empty() {
1417 Ok(None)
1418 } else {
1419 Ok(Some(arrow::compute::concat_batches(
1420 &batches[0].schema(),
1421 &batches,
1422 )?))
1423 }
1424 }
1425
1426 pub async fn scan_delta_table(
1429 &self,
1430 edge_type: &str,
1431 direction: &str,
1432 columns: &[&str],
1433 additional_filter: Option<&FilterExpr>,
1434 ) -> Result<Option<arrow_array::RecordBatch>> {
1435 let edge_hwm = self.snapshot_version_hwm();
1441 let backend = self.backend();
1442 let table_name = table_names::delta_table_name(edge_type, direction);
1443
1444 if !backend.table_exists(&table_name).await? {
1445 return Ok(None);
1446 }
1447
1448 let actual_columns =
1450 if let Some(table_schema) = backend.get_table_schema(&table_name).await? {
1451 let table_field_names: HashSet<&str> = table_schema
1452 .fields()
1453 .iter()
1454 .map(|f| f.name().as_str())
1455 .collect();
1456 columns
1457 .iter()
1458 .copied()
1459 .filter(|c| table_field_names.contains(c))
1460 .map(|s| s.to_string())
1461 .collect::<Vec<_>>()
1462 } else {
1463 return Ok(None);
1464 };
1465
1466 let filter = match (edge_hwm, additional_filter) {
1467 (Some(hwm), Some(f)) => Some(FilterExpr::all([
1468 FilterExpr::version_at_most(hwm),
1469 f.clone(),
1470 ])),
1471 (Some(hwm), None) => Some(FilterExpr::version_at_most(hwm)),
1472 (None, Some(f)) => Some(f.clone()),
1473 (None, None) => None,
1474 };
1475
1476 let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1477 if let Some(f) = filter {
1478 request = request.with_filter(f);
1479 }
1480
1481 let batches = backend.scan(request).await?;
1488 if batches.is_empty() {
1489 Ok(None)
1490 } else {
1491 Ok(Some(arrow::compute::concat_batches(
1492 &batches[0].schema(),
1493 &batches,
1494 )?))
1495 }
1496 }
1497
1498 pub async fn scan_main_vertex_table(
1503 &self,
1504 columns: &[&str],
1505 filter: Option<&FilterExpr>,
1506 ) -> Result<Option<arrow_array::RecordBatch>> {
1507 let backend = self.backend();
1508 let table_name = table_names::main_vertex_table_name();
1509
1510 if !backend.table_exists(table_name).await? {
1511 return Ok(None);
1512 }
1513
1514 let full_filter = match (self.version_high_water_mark(), filter) {
1516 (Some(hwm), Some(f)) => Some(FilterExpr::all([
1517 FilterExpr::version_at_most(hwm),
1518 f.clone(),
1519 ])),
1520 (Some(hwm), None) => Some(FilterExpr::version_at_most(hwm)),
1521 (None, Some(f)) => Some(f.clone()),
1522 (None, None) => None,
1523 };
1524
1525 let request = ScanRequest::all(table_name)
1526 .with_columns(columns.iter().map(|s| s.to_string()).collect());
1527 let request = match full_filter {
1528 Some(f) => request.with_filter(f),
1529 None => request,
1530 };
1531
1532 let batches = backend.scan(request).await?;
1539 if batches.is_empty() {
1540 Ok(None)
1541 } else {
1542 Ok(Some(arrow::compute::concat_batches(
1543 &batches[0].schema(),
1544 &batches,
1545 )?))
1546 }
1547 }
1548
1549 pub async fn scan_main_edge_table_stream(
1551 &self,
1552 filter: Option<&FilterExpr>,
1553 ) -> Result<
1554 Option<
1555 std::pin::Pin<Box<dyn futures::Stream<Item = Result<arrow_array::RecordBatch>> + Send>>,
1556 >,
1557 > {
1558 let backend = self.backend();
1559 let table_name = table_names::main_edge_table_name();
1560
1561 if !backend.table_exists(table_name).await? {
1562 return Ok(None);
1563 }
1564
1565 let mut request = ScanRequest::all(table_name);
1566 if let Some(f) = filter {
1567 request = request.with_filter(f.clone());
1568 }
1569
1570 let stream = backend.scan_stream(request).await?;
1571 Ok(Some(stream))
1572 }
1573
1574 pub async fn scan_vertex_table_stream(
1576 &self,
1577 label: &str,
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::vertex_table_name(label);
1585
1586 if !backend.table_exists(&table_name).await? {
1587 return Ok(None);
1588 }
1589
1590 let stream = backend.scan_stream(ScanRequest::all(&table_name)).await?;
1591 Ok(Some(stream))
1592 }
1593
1594 pub async fn find_vertex_by_ext_id(&self, ext_id: &str) -> Result<Option<Vid>> {
1596 MainVertexDataset::find_by_ext_id(self.backend(), ext_id, self.version_high_water_mark())
1597 .await
1598 }
1599
1600 pub async fn get_vertex_ext_ids(&self) -> Result<std::collections::HashMap<Vid, String>> {
1612 use arrow_array::StringArray;
1613 let backend = self.backend.as_ref();
1614 let vtable = table_names::main_vertex_table_name();
1615 let mut out = std::collections::HashMap::new();
1616 if !backend.table_exists(vtable).await? {
1617 return Ok(out);
1618 }
1619 let request = ScanRequest::all(vtable)
1620 .with_filter(FilterExpr::not_deleted())
1621 .with_columns(vec!["_vid".to_string(), "ext_id".to_string()]);
1622 let batches = backend
1623 .scan(request)
1624 .await
1625 .map_err(|e| anyhow!("get_vertex_ext_ids: {}", e))?;
1626 for batch in batches {
1627 let vids = batch
1628 .column_by_name("_vid")
1629 .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
1630 .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid _vid column"))?;
1631 let exts = batch
1632 .column_by_name("ext_id")
1633 .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1634 .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid ext_id column"))?;
1635 for i in 0..batch.num_rows() {
1636 if exts.is_null(i) {
1637 continue;
1638 }
1639 let ext = exts.value(i);
1640 if !ext.is_empty() {
1641 out.insert(Vid::from(vids.value(i)), ext.to_string());
1642 }
1643 }
1644 }
1645 Ok(out)
1646 }
1647
1648 pub async fn find_vertex_labels_by_vid(&self, vid: Vid) -> Result<Option<Vec<String>>> {
1650 MainVertexDataset::find_labels_by_vid(self.backend(), vid, self.version_high_water_mark())
1651 .await
1652 }
1653
1654 pub async fn find_edges_by_type_names(
1657 &self,
1658 type_names: &[&str],
1659 endpoint_filter: Option<(crate::storage::main_edge::EndpointSide, &[Vid])>,
1660 ) -> Result<Vec<(Eid, Vid, Vid, String, uni_common::Properties)>> {
1661 MainEdgeDataset::find_edges_by_type_names(self.backend(), type_names, endpoint_filter).await
1662 }
1663
1664 pub async fn scan_vertex_candidates(
1666 &self,
1667 label: &str,
1668 filter: Option<&FilterExpr>,
1669 ) -> Result<Vec<Vid>> {
1670 let backend = self.backend();
1671 let table_name = table_names::vertex_table_name(label);
1672
1673 if !backend.table_exists(&table_name).await? {
1674 return Ok(Vec::new());
1675 }
1676
1677 let full_filter = match filter {
1678 Some(f) => FilterExpr::all([FilterExpr::not_deleted(), f.clone()]),
1679 None => FilterExpr::not_deleted(),
1680 };
1681
1682 let request = ScanRequest::all(&table_name)
1683 .with_filter(full_filter)
1684 .with_columns(vec!["_vid".to_string()]);
1685
1686 let batches = backend.scan(request).await?;
1687
1688 let mut vids = Vec::new();
1689 for batch in batches {
1690 let vid_col = batch
1691 .column_by_name("_vid")
1692 .ok_or(anyhow!("Missing _vid"))?
1693 .as_any()
1694 .downcast_ref::<UInt64Array>()
1695 .ok_or(anyhow!("Invalid _vid"))?;
1696 for i in 0..batch.num_rows() {
1697 vids.push(Vid::from(vid_col.value(i)));
1698 }
1699 }
1700 Ok(vids)
1701 }
1702
1703 pub fn vertex_dataset(&self, label: &str) -> Result<VertexDataset> {
1709 let schema = self.schema_manager.schema();
1710 let label_meta = schema
1711 .labels
1712 .get(label)
1713 .ok_or_else(|| anyhow!("Label '{}' not found", label))?;
1714 Ok(VertexDataset::new(&self.base_uri, label, label_meta.id))
1715 }
1716
1717 #[cfg(feature = "lance-backend")]
1718 pub fn edge_dataset(
1719 &self,
1720 edge_type: &str,
1721 src_label: &str,
1722 dst_label: &str,
1723 ) -> Result<EdgeDataset> {
1724 let key = format!("edges_{edge_type}");
1725 match self.fork_branch_for(&key) {
1726 Some(branch) => Ok(EdgeDataset::new_branched(
1727 &self.base_uri,
1728 edge_type,
1729 src_label,
1730 dst_label,
1731 branch,
1732 )),
1733 None => Ok(EdgeDataset::new(
1734 &self.base_uri,
1735 edge_type,
1736 src_label,
1737 dst_label,
1738 )),
1739 }
1740 }
1741
1742 pub fn delta_dataset(&self, edge_type: &str, direction: &str) -> Result<DeltaDataset> {
1743 let key = format!("deltas_{edge_type}_{direction}");
1744 match self.fork_branch_for(&key) {
1745 Some(branch) => Ok(DeltaDataset::new_branched(
1746 &self.base_uri,
1747 edge_type,
1748 direction,
1749 branch,
1750 )),
1751 None => Ok(DeltaDataset::new(&self.base_uri, edge_type, direction)),
1752 }
1753 }
1754
1755 pub fn adjacency_dataset(
1756 &self,
1757 edge_type: &str,
1758 label: &str,
1759 direction: &str,
1760 ) -> Result<AdjacencyDataset> {
1761 let key = crate::backend::table_names::adjacency_table_name(edge_type, direction);
1768 match self.fork_branch_for(&key) {
1769 Some(branch) => Ok(AdjacencyDataset::new_branched(
1770 &self.base_uri,
1771 edge_type,
1772 label,
1773 direction,
1774 branch,
1775 )),
1776 None => Ok(AdjacencyDataset::new(
1777 &self.base_uri,
1778 edge_type,
1779 label,
1780 direction,
1781 )),
1782 }
1783 }
1784
1785 fn fork_branch_for(&self, dataset_name: &str) -> Option<String> {
1789 self.fork_scope
1790 .as_ref()
1791 .and_then(|s| s.branch_for(dataset_name))
1792 }
1793
1794 pub fn main_vertex_dataset(&self) -> MainVertexDataset {
1799 MainVertexDataset::new(&self.base_uri)
1800 }
1801
1802 pub fn main_edge_dataset(&self) -> MainEdgeDataset {
1807 MainEdgeDataset::new(&self.base_uri)
1808 }
1809
1810 #[cfg(feature = "lance-backend")]
1811 pub fn uid_index(&self, label: &str) -> Result<UidIndex> {
1812 Ok(UidIndex::new(&self.base_uri, label))
1813 }
1814
1815 #[cfg(feature = "lance-backend")]
1816 pub async fn inverted_index(&self, label: &str, property: &str) -> Result<InvertedIndex> {
1817 let schema = self.schema_manager.schema();
1818 let config = schema
1819 .indexes
1820 .iter()
1821 .find_map(|idx| match idx {
1822 IndexDefinition::Inverted(cfg)
1823 if cfg.label == label && cfg.property == property =>
1824 {
1825 Some(cfg.clone())
1826 }
1827 _ => None,
1828 })
1829 .ok_or_else(|| anyhow!("Inverted index not found for {}.{}", label, property))?;
1830
1831 InvertedIndex::new(&self.base_uri, config).await
1832 }
1833
1834 #[expect(clippy::too_many_arguments)]
1835 pub async fn vector_search(
1836 &self,
1837 label: &str,
1838 property: &str,
1839 query: &[f32],
1840 k: usize,
1841 filter: Option<&str>,
1842 opts: VectorQueryOpts,
1843 ctx: Option<&QueryContext>,
1844 ) -> Result<Vec<(Vid, f32)>> {
1845 use crate::backend::types::DistanceMetric as BackendMetric;
1846
1847 let schema = self.schema_manager.schema();
1849
1850 if let Some(meta) = schema
1855 .properties
1856 .get(label)
1857 .and_then(|props| props.get(property))
1858 && let uni_common::core::schema::DataType::Vector { dimensions } = &meta.r#type
1859 && query.len() != *dimensions
1860 {
1861 return Err(anyhow!(
1862 "vector dimension mismatch: query vector has {} dimensions but '{}.{}' is \
1863 declared VECTOR({})",
1864 query.len(),
1865 label,
1866 property,
1867 dimensions
1868 ));
1869 }
1870
1871 let metric = schema
1872 .vector_index_for_property(label, property)
1873 .map(|config| config.metric.clone())
1874 .unwrap_or(DistanceMetric::L2);
1875
1876 let backend = self.backend.as_ref();
1877 let name = table_names::vertex_table_name(label);
1878
1879 let mut results = Vec::new();
1880
1881 if backend.table_exists(&name).await? {
1883 let backend_metric = match &metric {
1884 DistanceMetric::L2 => BackendMetric::L2,
1885 DistanceMetric::Cosine => BackendMetric::Cosine,
1886 DistanceMetric::Dot => BackendMetric::Dot,
1887 DistanceMetric::L1 => BackendMetric::L2,
1892 _ => BackendMetric::L2,
1893 };
1894
1895 let mut filter_parts = vec![Self::build_active_filter(filter)];
1897 if ctx.is_some()
1898 && let Some(hwm) = self.version_high_water_mark()
1899 {
1900 filter_parts.push(FilterExpr::version_at_most(hwm));
1901 }
1902 let combined_filter = FilterExpr::all(filter_parts);
1903
1904 let fetch_k = if matches!(metric, DistanceMetric::L1) {
1908 backend.count_rows(&name, None).await.unwrap_or(k).max(k)
1909 } else {
1910 k
1911 };
1912
1913 let batches = backend
1914 .vector_search(
1915 &name,
1916 property,
1917 query,
1918 fetch_k,
1919 backend_metric,
1920 combined_filter,
1921 opts,
1922 )
1923 .await?;
1924
1925 results = extract_vid_and_vector_pairs(&batches, "_vid", property)?
1934 .into_iter()
1935 .filter(|(_, emb)| emb.len() == query.len())
1936 .map(|(vid, emb)| (vid, metric.compute_distance(&emb, query)))
1937 .collect();
1938 }
1939
1940 if let Some(qctx) = ctx {
1942 merge_l0_into_vector_results(&mut results, qctx, label, property, query, k, &metric);
1943 }
1944
1945 results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1950 results.truncate(k);
1951
1952 Ok(results)
1953 }
1954
1955 #[expect(clippy::too_many_arguments)]
1966 pub async fn muvera_fde_candidates(
1967 &self,
1968 label: &str,
1969 fde_column: &str,
1970 fde_query: &[f32],
1971 k: usize,
1972 filter: Option<&str>,
1973 opts: VectorQueryOpts,
1974 ctx: Option<&QueryContext>,
1975 ) -> Result<Vec<(Vid, f32)>> {
1976 use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
1977
1978 let backend = self.backend.as_ref();
1979 let name = table_names::vertex_table_name(label);
1980 if !backend.table_exists(&name).await? {
1984 return Ok(Vec::new());
1985 }
1986
1987 let mut filter_parts = vec![Self::build_active_filter(filter)];
1988 if ctx.is_some()
1989 && let Some(hwm) = self.version_high_water_mark()
1990 {
1991 filter_parts.push(FilterExpr::version_at_most(hwm));
1992 }
1993 let combined_filter = FilterExpr::all(filter_parts);
1994
1995 let batches = backend
1996 .vector_search(
1997 &name,
1998 fde_column,
1999 fde_query,
2000 k,
2001 BackendMetric::Dot,
2002 combined_filter,
2003 opts,
2004 )
2005 .await?;
2006 extract_vid_score_pairs(&batches, "_vid", "_distance")
2007 }
2008
2009 #[expect(clippy::too_many_arguments)]
2030 pub async fn multivector_search(
2031 &self,
2032 label: &str,
2033 property: &str,
2034 query: &[Vec<f32>],
2035 k: usize,
2036 filter: Option<&str>,
2037 opts: VectorQueryOpts,
2038 ctx: Option<&QueryContext>,
2039 ) -> Result<Vec<(Vid, f32)>> {
2040 use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
2041
2042 let schema = self.schema_manager.schema();
2043 let metric = schema
2044 .vector_index_for_property(label, property)
2045 .map(|config| config.metric.clone())
2046 .unwrap_or(DistanceMetric::Cosine);
2047
2048 let backend = self.backend.as_ref();
2049 let name = table_names::vertex_table_name(label);
2050
2051 let branched = self
2062 .fork_scope
2063 .as_ref()
2064 .is_some_and(|s| s.branch_for(&name).is_some());
2065 if branched {
2066 if !backend.table_exists(&name).await? {
2069 return Ok(Vec::new());
2070 }
2071 let mut filter_parts = vec![Self::build_active_filter(filter)];
2072 if ctx.is_some()
2073 && let Some(hwm) = self.version_high_water_mark()
2074 {
2075 filter_parts.push(FilterExpr::version_at_most(hwm));
2076 }
2077 let request = ScanRequest::all(&name)
2078 .with_filter(FilterExpr::all(filter_parts))
2079 .with_columns(vec!["_vid".to_string()]);
2080 let batches = backend.scan(request).await?;
2081 let mut results = Vec::new();
2082 for batch in batches {
2083 let vid_col = batch
2084 .column_by_name("_vid")
2085 .ok_or(anyhow!("Missing _vid"))?
2086 .as_any()
2087 .downcast_ref::<UInt64Array>()
2088 .ok_or(anyhow!("Invalid _vid"))?;
2089 for i in 0..batch.num_rows() {
2090 results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2091 }
2092 }
2093 return Ok(results);
2094 }
2095
2096 let mut results = Vec::new();
2097 if backend.table_exists(&name).await? {
2099 let backend_metric = match &metric {
2100 DistanceMetric::L2 => BackendMetric::L2,
2101 DistanceMetric::Cosine => BackendMetric::Cosine,
2102 DistanceMetric::Dot => BackendMetric::Dot,
2103 _ => BackendMetric::Cosine,
2104 };
2105
2106 let mut filter_parts = vec![Self::build_active_filter(filter)];
2107 if ctx.is_some()
2108 && let Some(hwm) = self.version_high_water_mark()
2109 {
2110 filter_parts.push(FilterExpr::version_at_most(hwm));
2111 }
2112 let combined_filter = FilterExpr::all(filter_parts);
2113
2114 let batches = backend
2115 .multivector_search(
2116 &name,
2117 property,
2118 query,
2119 k,
2120 backend_metric,
2121 combined_filter,
2122 opts,
2123 )
2124 .await?;
2125 results = extract_vid_score_pairs(&batches, "_vid", "_distance")?;
2126 }
2127
2128 Ok(results)
2129 }
2130
2131 pub async fn sparse_search(
2144 &self,
2145 label: &str,
2146 property: &str,
2147 query: &[(u32, f32)],
2148 k: usize,
2149 ) -> Result<Vec<(Vid, f32)>> {
2150 #[cfg(feature = "lance-backend")]
2151 {
2152 let name = table_names::vertex_table_name(label);
2153 let branched = self
2154 .fork_scope
2155 .as_ref()
2156 .is_some_and(|s| s.branch_for(&name).is_some());
2157 if branched {
2158 let backend = self.backend.as_ref();
2180 if !backend.table_exists(&name).await? {
2183 return Ok(Vec::new());
2184 }
2185 let request = ScanRequest::all(&name)
2186 .with_filter(Self::build_active_filter(None))
2187 .with_columns(vec!["_vid".to_string()]);
2188 let batches = backend.scan(request).await?;
2189 let mut results = Vec::new();
2190 for batch in batches {
2191 let vid_col = batch
2192 .column_by_name("_vid")
2193 .ok_or_else(|| anyhow!("Missing _vid"))?
2194 .as_any()
2195 .downcast_ref::<UInt64Array>()
2196 .ok_or_else(|| anyhow!("Invalid _vid"))?;
2197 for i in 0..batch.num_rows() {
2198 results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2199 }
2200 }
2201 return Ok(results);
2202 }
2203 match self
2204 .index_manager()
2205 .sparse_vector_index(label, property)
2206 .await
2207 {
2208 Ok(idx) => idx.query_topk(query, k).await,
2209 Err(_) => Ok(Vec::new()),
2210 }
2211 }
2212 #[cfg(not(feature = "lance-backend"))]
2213 {
2214 let _ = (label, property, query, k);
2215 Ok(Vec::new())
2216 }
2217 }
2218
2219 pub async fn fts_search(
2235 &self,
2236 label: &str,
2237 property: &str,
2238 query: &str,
2239 k: usize,
2240 filter: Option<&str>,
2241 ctx: Option<&QueryContext>,
2242 ) -> Result<Vec<(Vid, f32)>> {
2243 let backend = self.backend.as_ref();
2244 let name = table_names::vertex_table_name(label);
2245
2246 let mut results = if backend.table_exists(&name).await? {
2247 let mut filter_parts = vec![Self::build_active_filter(filter)];
2249 if ctx.is_some()
2250 && let Some(hwm) = self.version_high_water_mark()
2251 {
2252 filter_parts.push(FilterExpr::version_at_most(hwm));
2253 }
2254 let combined_filter = FilterExpr::all(filter_parts);
2255
2256 let batches = backend
2257 .full_text_search(&name, property, query, k, combined_filter)
2258 .await?;
2259
2260 let mut fts_results = extract_vid_score_pairs(&batches, "_vid", "_score")?;
2261 fts_results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2263 fts_results
2264 } else {
2265 Vec::new()
2266 };
2267
2268 if let Some(qctx) = ctx {
2270 merge_l0_into_fts_results(&mut results, qctx, label, property, query, k);
2271 }
2272
2273 Ok(results)
2274 }
2275
2276 #[cfg(feature = "lance-backend")]
2277 pub async fn get_vertex_by_uid(&self, uid: &UniId, label: &str) -> Result<Option<Vid>> {
2278 let index = self.uid_index(label)?;
2279 index.get_vid(uid).await
2280 }
2281
2282 #[cfg(feature = "lance-backend")]
2283 pub async fn insert_vertex_with_uid(&self, label: &str, vid: Vid, uid: UniId) -> Result<()> {
2284 let index = self.uid_index(label)?;
2285 index.write_mapping(&[(uid, vid)]).await
2286 }
2287
2288 pub async fn load_subgraph(
2289 &self,
2290 start_vids: &[Vid],
2291 edge_types: &[u32],
2292 max_hops: usize,
2293 direction: GraphDirection,
2294 l0: Option<&L0Buffer>,
2295 ) -> Result<WorkingGraph> {
2296 let mut graph = WorkingGraph::new();
2297 let schema = self.schema_manager.schema();
2298
2299 let label_map: HashMap<u16, String> = schema
2301 .labels
2302 .values()
2303 .map(|meta| {
2304 (
2305 meta.id,
2306 schema.label_name_by_id(meta.id).unwrap().to_owned(),
2307 )
2308 })
2309 .collect();
2310
2311 let edge_type_map: HashMap<u32, String> = schema
2312 .edge_types
2313 .values()
2314 .map(|meta| {
2315 (
2316 meta.id,
2317 schema.edge_type_name_by_id(meta.id).unwrap().to_owned(),
2318 )
2319 })
2320 .collect();
2321
2322 let target_edge_types: HashSet<u32> = edge_types.iter().copied().collect();
2323
2324 let mut frontier: Vec<Vid> = start_vids.to_vec();
2326 let mut visited: HashSet<Vid> = HashSet::new();
2327
2328 for &vid in start_vids {
2330 graph.add_vertex(vid);
2331 }
2332
2333 for _hop in 0..max_hops {
2334 let mut next_frontier = HashSet::new();
2335
2336 for &vid in &frontier {
2337 if visited.contains(&vid) {
2338 continue;
2339 }
2340 visited.insert(vid);
2341 graph.add_vertex(vid);
2342
2343 for &etype_id in &target_edge_types {
2345 let etype_name = edge_type_map
2346 .get(&etype_id)
2347 .ok_or_else(|| anyhow!("Unknown edge type ID: {}", etype_id))?;
2348
2349 let (dir_str, neighbor_is_dst) = match direction {
2353 GraphDirection::Outgoing => ("fwd", true),
2354 GraphDirection::Incoming => ("bwd", false),
2355 };
2356
2357 let mut edges: HashMap<Eid, EdgeState> = HashMap::new();
2358
2359 let _edge_ver = self
2364 .pinned_snapshot
2365 .as_ref()
2366 .and_then(|s| s.edges.get(etype_name).map(|es| es.lance_version));
2367
2368 let backend = self.backend();
2370 for current_src_label in label_map.values() {
2371 let adj_ds =
2372 match self.adjacency_dataset(etype_name, current_src_label, dir_str) {
2373 Ok(ds) => ds,
2374 Err(_) => continue,
2375 };
2376 if let Some((neighbors, eids)) =
2377 adj_ds.read_adjacency_backend(backend, vid).await?
2378 {
2379 for (n, eid) in neighbors.into_iter().zip(eids) {
2380 edges.insert(
2381 eid,
2382 EdgeState {
2383 neighbor: n,
2384 version: 0,
2385 deleted: false,
2386 },
2387 );
2388 }
2389 break; }
2391 }
2392
2393 let delta_ds = self.delta_dataset(etype_name, dir_str)?;
2395 let delta_entries = delta_ds
2396 .read_deltas(backend, vid, &schema, self.snapshot_version_hwm())
2397 .await?;
2398 Self::apply_delta_to_edges(&mut edges, delta_entries, neighbor_is_dst);
2399
2400 if let Some(l0) = l0 {
2402 Self::apply_l0_to_edges(&mut edges, l0, vid, etype_id, direction);
2403 }
2404
2405 Self::add_edges_to_graph(
2407 &mut graph,
2408 edges,
2409 vid,
2410 etype_id,
2411 neighbor_is_dst,
2412 &visited,
2413 &mut next_frontier,
2414 );
2415 }
2416 }
2417 frontier = next_frontier.into_iter().collect();
2418
2419 if frontier.is_empty() {
2421 break;
2422 }
2423 }
2424
2425 Ok(graph)
2426 }
2427
2428 fn apply_delta_to_edges(
2430 edges: &mut HashMap<Eid, EdgeState>,
2431 delta_entries: Vec<crate::storage::delta::L1Entry>,
2432 neighbor_is_dst: bool,
2433 ) {
2434 for entry in delta_entries {
2435 let neighbor = if neighbor_is_dst {
2436 entry.dst_vid
2437 } else {
2438 entry.src_vid
2439 };
2440 let current_ver = edges.get(&entry.eid).map(|s| s.version).unwrap_or(0);
2441
2442 if entry.version > current_ver {
2443 edges.insert(
2444 entry.eid,
2445 EdgeState {
2446 neighbor,
2447 version: entry.version,
2448 deleted: matches!(entry.op, Op::Delete),
2449 },
2450 );
2451 }
2452 }
2453 }
2454
2455 fn apply_l0_to_edges(
2457 edges: &mut HashMap<Eid, EdgeState>,
2458 l0: &L0Buffer,
2459 vid: Vid,
2460 etype_id: u32,
2461 direction: GraphDirection,
2462 ) {
2463 let l0_neighbors = l0.get_neighbors(vid, etype_id, direction);
2464 for (neighbor, eid, ver) in l0_neighbors {
2465 let current_ver = edges.get(&eid).map(|s| s.version).unwrap_or(0);
2466 if ver > current_ver {
2467 edges.insert(
2468 eid,
2469 EdgeState {
2470 neighbor,
2471 version: ver,
2472 deleted: false,
2473 },
2474 );
2475 }
2476 }
2477
2478 for (eid, state) in edges.iter_mut() {
2480 if l0.is_tombstoned(*eid) {
2481 state.deleted = true;
2482 }
2483 }
2484 }
2485
2486 fn add_edges_to_graph(
2488 graph: &mut WorkingGraph,
2489 edges: HashMap<Eid, EdgeState>,
2490 vid: Vid,
2491 etype_id: u32,
2492 neighbor_is_dst: bool,
2493 visited: &HashSet<Vid>,
2494 next_frontier: &mut HashSet<Vid>,
2495 ) {
2496 for (eid, state) in edges {
2497 if state.deleted {
2498 continue;
2499 }
2500 graph.add_vertex(state.neighbor);
2501
2502 if !visited.contains(&state.neighbor) {
2503 next_frontier.insert(state.neighbor);
2504 }
2505
2506 if neighbor_is_dst {
2507 graph.add_edge(vid, state.neighbor, eid, etype_id);
2508 } else {
2509 graph.add_edge(state.neighbor, vid, eid, etype_id);
2510 }
2511 }
2512 }
2513}
2514
2515fn extract_vid_score_pairs(
2517 batches: &[arrow_array::RecordBatch],
2518 vid_column: &str,
2519 score_column: &str,
2520) -> Result<Vec<(Vid, f32)>> {
2521 let mut results = Vec::new();
2522 for batch in batches {
2523 let vid_col = batch
2524 .column_by_name(vid_column)
2525 .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2526 .as_any()
2527 .downcast_ref::<UInt64Array>()
2528 .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2529
2530 let score_col = batch
2531 .column_by_name(score_column)
2532 .ok_or_else(|| anyhow!("Missing {} column", score_column))?
2533 .as_any()
2534 .downcast_ref::<Float32Array>()
2535 .ok_or_else(|| anyhow!("Invalid {} column type", score_column))?;
2536
2537 for i in 0..batch.num_rows() {
2538 results.push((Vid::from(vid_col.value(i)), score_col.value(i)));
2539 }
2540 }
2541 Ok(results)
2542}
2543
2544fn extract_vid_and_vector_pairs(
2553 batches: &[arrow_array::RecordBatch],
2554 vid_column: &str,
2555 vector_column: &str,
2556) -> Result<Vec<(Vid, Vec<f32>)>> {
2557 use arrow_array::{Array, FixedSizeListArray};
2558 let mut out = Vec::new();
2559 for batch in batches {
2560 let vid_col = batch
2561 .column_by_name(vid_column)
2562 .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2563 .as_any()
2564 .downcast_ref::<UInt64Array>()
2565 .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2566
2567 let vec_col = batch
2568 .column_by_name(vector_column)
2569 .ok_or_else(|| anyhow!("Missing vector column {}", vector_column))?
2570 .as_any()
2571 .downcast_ref::<FixedSizeListArray>()
2572 .ok_or_else(|| anyhow!("Vector column {} is not FixedSizeList", vector_column))?;
2573
2574 for i in 0..batch.num_rows() {
2575 if vec_col.is_null(i) {
2576 continue;
2577 }
2578 let element = vec_col.value(i);
2579 let floats = element
2580 .as_any()
2581 .downcast_ref::<Float32Array>()
2582 .ok_or_else(|| {
2583 anyhow!("Vector column {} inner type is not Float32", vector_column)
2584 })?;
2585 let emb: Vec<f32> = (0..floats.len()).map(|j| floats.value(j)).collect();
2586 out.push((Vid::from(vid_col.value(i)), emb));
2587 }
2588 }
2589 Ok(out)
2590}
2591
2592fn extract_embedding_from_props(
2606 props: &uni_common::Properties,
2607 property: &str,
2608) -> Option<Vec<f32>> {
2609 Vec::<f32>::try_from(props.get(property)?).ok()
2610}
2611
2612fn merge_l0_into_vector_results(
2622 results: &mut Vec<(Vid, f32)>,
2623 ctx: &QueryContext,
2624 label: &str,
2625 property: &str,
2626 query: &[f32],
2627 k: usize,
2628 metric: &DistanceMetric,
2629) {
2630 let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2632 ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2633 buffers.push(Arc::clone(&ctx.l0));
2634 if let Some(ref txn) = ctx.transaction_l0 {
2635 buffers.push(Arc::clone(txn));
2636 }
2637
2638 let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2640 let mut tombstoned: HashSet<Vid> = HashSet::new();
2642
2643 for buf_arc in &buffers {
2644 let buf = buf_arc.read();
2645
2646 for &vid in &buf.vertex_tombstones {
2648 tombstoned.insert(vid);
2649 }
2650
2651 for (&vid, labels) in &buf.vertex_labels {
2653 if !labels.iter().any(|l| l == label) {
2654 continue;
2655 }
2656 if let Some(props) = buf.vertex_properties.get(&vid)
2657 && let Some(emb) = extract_embedding_from_props(props, property)
2658 {
2659 if emb.len() != query.len() {
2660 continue; }
2668 let dist = metric.compute_distance(&emb, query);
2669 l0_candidates.insert(vid, dist);
2671 tombstoned.remove(&vid);
2673 }
2674 }
2675 }
2676
2677 if l0_candidates.is_empty() && tombstoned.is_empty() {
2679 return;
2680 }
2681
2682 results.retain(|(vid, _)| !tombstoned.contains(vid));
2684
2685 for (vid, dist) in &l0_candidates {
2687 if tombstoned.contains(vid) {
2693 continue;
2694 }
2695 if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2696 existing.1 = *dist;
2697 } else {
2698 results.push((*vid, *dist));
2699 }
2700 }
2701
2702 results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
2704 results.truncate(k);
2705}
2706
2707pub fn collect_l0_label_candidates(ctx: &QueryContext, label: &str) -> (Vec<Vid>, HashSet<Vid>) {
2721 let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2723 ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2724 buffers.push(Arc::clone(&ctx.l0));
2725 if let Some(ref txn) = ctx.transaction_l0 {
2726 buffers.push(Arc::clone(txn));
2727 }
2728
2729 let mut live: HashSet<Vid> = HashSet::new();
2730 let mut tombstoned: HashSet<Vid> = HashSet::new();
2731
2732 for buf_arc in &buffers {
2733 let buf = buf_arc.read();
2734
2735 for &vid in &buf.vertex_tombstones {
2737 tombstoned.insert(vid);
2738 live.remove(&vid);
2739 }
2740
2741 for (&vid, labels) in &buf.vertex_labels {
2744 if !labels.iter().any(|l| l == label) {
2745 continue;
2746 }
2747 if buf.vertex_properties.contains_key(&vid) {
2748 live.insert(vid);
2749 tombstoned.remove(&vid);
2750 }
2751 }
2752 }
2753
2754 (live.into_iter().collect(), tombstoned)
2755}
2756
2757fn compute_text_relevance(query: &str, text: &str) -> f32 {
2762 let query_tokens: HashSet<String> =
2763 query.split_whitespace().map(|t| t.to_lowercase()).collect();
2764 if query_tokens.is_empty() {
2765 return 0.0;
2766 }
2767 let text_tokens: HashSet<String> = text.split_whitespace().map(|t| t.to_lowercase()).collect();
2768 let hits = query_tokens
2769 .iter()
2770 .filter(|t| text_tokens.contains(t.as_str()))
2771 .count();
2772 hits as f32 / query_tokens.len() as f32
2773}
2774
2775fn extract_text_from_props<'a>(
2777 props: &'a uni_common::Properties,
2778 property: &str,
2779) -> Option<&'a str> {
2780 props.get(property)?.as_str()
2781}
2782
2783fn merge_l0_into_fts_results(
2793 results: &mut Vec<(Vid, f32)>,
2794 ctx: &QueryContext,
2795 label: &str,
2796 property: &str,
2797 query: &str,
2798 k: usize,
2799) {
2800 let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2802 ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2803 buffers.push(Arc::clone(&ctx.l0));
2804 if let Some(ref txn) = ctx.transaction_l0 {
2805 buffers.push(Arc::clone(txn));
2806 }
2807
2808 let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2810 let mut tombstoned: HashSet<Vid> = HashSet::new();
2812
2813 for buf_arc in &buffers {
2814 let buf = buf_arc.read();
2815
2816 for &vid in &buf.vertex_tombstones {
2818 tombstoned.insert(vid);
2819 }
2820
2821 for (&vid, labels) in &buf.vertex_labels {
2823 if !labels.iter().any(|l| l == label) {
2824 continue;
2825 }
2826 if let Some(props) = buf.vertex_properties.get(&vid)
2827 && let Some(text) = extract_text_from_props(props, property)
2828 {
2829 let score = compute_text_relevance(query, text);
2830 if score > 0.0 {
2831 l0_candidates.insert(vid, score);
2833 }
2834 tombstoned.remove(&vid);
2836 }
2837 }
2838 }
2839
2840 if l0_candidates.is_empty() && tombstoned.is_empty() {
2842 return;
2843 }
2844
2845 results.retain(|(vid, _)| !tombstoned.contains(vid));
2847
2848 for (vid, score) in &l0_candidates {
2850 if tombstoned.contains(vid) {
2854 continue;
2855 }
2856 if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2857 existing.1 = *score;
2858 } else {
2859 results.push((*vid, *score));
2860 }
2861 }
2862
2863 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2865 results.truncate(k);
2866}