Skip to main content

uni_store/storage/
manager.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4use crate::backend::StorageBackend;
5#[cfg(feature = "lance-backend")]
6use crate::backend::lance::LanceDbBackend;
7use crate::backend::table_names;
8use crate::backend::types::{ScanRequest, VectorQueryOpts};
9use crate::compaction::{CompactionStats, CompactionStatus, CompactionTask};
10use crate::runtime::WorkingGraph;
11use crate::runtime::context::QueryContext;
12use crate::runtime::l0::L0Buffer;
13use crate::storage::adjacency::AdjacencyDataset;
14use crate::storage::compaction::Compactor;
15use crate::storage::delta::{DeltaDataset, ENTRY_SIZE_ESTIMATE, Op};
16use crate::storage::direction::Direction;
17#[cfg(feature = "lance-backend")]
18use crate::storage::edge::EdgeDataset;
19#[cfg(feature = "lance-backend")]
20use crate::storage::index::UidIndex;
21#[cfg(feature = "lance-backend")]
22use crate::storage::inverted_index::InvertedIndex;
23use crate::storage::main_edge::MainEdgeDataset;
24use crate::storage::main_vertex::MainVertexDataset;
25use crate::storage::vertex::VertexDataset;
26use anyhow::{Result, anyhow};
27use arrow_array::{Array, Float32Array, TimestampNanosecondArray, UInt64Array};
28use object_store::ObjectStore;
29#[cfg(feature = "lance-backend")]
30use object_store::local::LocalFileSystem;
31use parking_lot::RwLock;
32use std::collections::{HashMap, HashSet};
33use std::sync::{Arc, Mutex};
34use std::time::{Duration, SystemTime, UNIX_EPOCH};
35use tracing::warn;
36use uni_common::config::UniConfig;
37#[cfg(feature = "lance-backend")]
38use uni_common::core::id::UniId;
39use uni_common::core::id::{Eid, Vid};
40#[cfg(feature = "lance-backend")]
41use uni_common::core::schema::IndexDefinition;
42use uni_common::core::schema::{DistanceMetric, SchemaManager};
43use uni_common::sync::acquire_mutex;
44
45use crate::snapshot::manager::SnapshotManager;
46use crate::storage::IndexManager;
47use crate::storage::adjacency_manager::AdjacencyManager;
48use crate::storage::resilient_store::ResilientObjectStore;
49
50use uni_common::core::snapshot::SnapshotManifest;
51
52use uni_common::graph::simple_graph::Direction as GraphDirection;
53
54/// Edge state during subgraph loading - tracks version and deletion status.
55struct 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    /// Counter of in-flight `flush_to_l1` operations. Compaction skips
70    /// delta-clear when this is non-zero to avoid wiping rows a flush is
71    /// about to append. Counter (not bool) so multiple async flushes can
72    /// be in flight concurrently.
73    pub flush_in_progress: std::sync::atomic::AtomicUsize,
74    /// Optional pinned snapshot for time-travel
75    pinned_snapshot: Option<SnapshotManifest>,
76    /// Optional row-version pin for transaction snapshot reads (C2).
77    ///
78    /// When set, L1 scans filter to `_version <= hwm` exactly like a pinned
79    /// snapshot, but WITHOUT a manifest: a read-write transaction pins the
80    /// version counter observed at begin (`SnapshotView.started_at_version`)
81    /// so an L0→L1 flush completing mid-transaction cannot leak
82    /// post-snapshot rows into its scans. Mutually exclusive with
83    /// `pinned_snapshot`.
84    pinned_version_hwm: Option<u64>,
85    /// Optional fork scope for branch-aware reads (Phase 1 read-only).
86    ///
87    /// Mutually exclusive with `pinned_snapshot`: a single
88    /// `StorageManager` is either pinned to a snapshot or scoped to a
89    /// fork, never both. Phase 4's `pin_to_version` on a forked session
90    /// builds a separate combined manager out of band; Phase 1 forbids
91    /// mixing.
92    fork_scope: Option<Arc<crate::fork::ForkScope>>,
93    /// Pluggable storage backend.
94    backend: Arc<dyn StorageBackend>,
95    /// In-memory VID-to-labels index for O(1) label lookups.
96    ///
97    /// Always present: populated at startup via [`Self::rebuild_vid_labels_index`]
98    /// and kept current at flush time. Traversal-time label predicates
99    /// (`MATCH (a)-[r]->(b:B)`) read it to resolve labels for vertices that
100    /// have aged out of L0 into Lance storage — notably on forks, whose data
101    /// is flushed to Lance before branching.
102    vid_labels_index: Arc<parking_lot::RwLock<crate::storage::vid_labels::VidLabelsIndex>>,
103    /// Optional plugin registry for registry-dispatched CRDT merges on the
104    /// durable paths (compaction, L0 flush).
105    ///
106    /// Behavior-preserving when absent: the durable merge helpers fall back to
107    /// [`uni_crdt::Crdt::try_merge`] bit-for-bit when no
108    /// [`uni_plugin::traits::crdt::CrdtKindProvider`] is registered. Threaded
109    /// down to the writer's `L0Manager` (and hence each `L0Buffer`) and read by
110    /// the [`crate::storage::compaction::Compactor`] via
111    /// [`Self::plugin_registry`].
112    plugin_registry: Option<Arc<uni_plugin::PluginRegistry>>,
113}
114
115/// RAII counter increment for `StorageManager.flush_in_progress`.
116///
117/// Acquired during the rotate phase of a flush (see
118/// `runtime::writer::flush_l0_rotate`) and dropped when the full
119/// rotate/stream/finalize pipeline completes. Compaction's delta-clear
120/// gate skips while this counter is non-zero, so the counter must
121/// reflect "flush has started, has not completed" — including any
122/// async stream phase running on a spawned task.
123pub struct FlushInProgressGuard {
124    storage: Arc<StorageManager>,
125}
126
127impl FlushInProgressGuard {
128    pub fn new(storage: &Arc<StorageManager>) -> Self {
129        storage
130            .flush_in_progress
131            .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
132        Self {
133            storage: storage.clone(),
134        }
135    }
136}
137
138impl Drop for FlushInProgressGuard {
139    fn drop(&mut self) {
140        // M-PANIC-IS-STOP: must not panic in Drop. Atomic op cannot fail.
141        self.storage
142            .flush_in_progress
143            .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
144    }
145}
146
147/// Whether a Lance error represents a commit conflict that retrying may
148/// resolve. These fire under async-flush when ≥2 streams concurrently try
149/// to create the same table OR when an Append races with a still-in-progress
150/// Overwrite (create_table). See the Lance commit-conflict-resolver in
151/// `lance-3.0.1/src/io/commit/conflict_resolver.rs`.
152fn is_lance_conflict(err: &anyhow::Error) -> bool {
153    let msg = err.to_string();
154    msg.contains("Incompatible transaction") || msg.contains("conflict")
155}
156
157/// Runs `op` with exponential-backoff retry on Lance commit conflicts.
158/// Up to 10 attempts (~10s worst case); backoff is 1ms, 2ms, 4ms, ...,
159/// 512ms. Non-conflict errors return immediately. `op` is re-invoked each
160/// attempt so it can re-check table existence and adjust strategy.
161async fn retry_on_lance_conflict<F, Fut>(mut op: F) -> anyhow::Result<()>
162where
163    F: FnMut() -> Fut,
164    Fut: std::future::Future<Output = anyhow::Result<()>>,
165{
166    for attempt in 0u32..10 {
167        match op().await {
168            Ok(()) => return Ok(()),
169            Err(e) => {
170                if !is_lance_conflict(&e) || attempt == 9 {
171                    return Err(e);
172                }
173                let backoff_ms = 1u64 << attempt;
174                tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
175            }
176        }
177    }
178    unreachable!("retry loop exits via Ok or Err")
179}
180
181/// MergeInsert sibling of `write_batch_with_lance_conflict_retry`.
182///
183/// Source `batch` must contain the join columns in `on` plus any
184/// columns to update. Matched rows have `WhenMatched::UpdateAll`
185/// applied; unmatched source rows are dropped (partial writes never
186/// INSERT). Returns an error if the target table does not exist.
187/// Retries on Lance commit conflicts via `retry_on_lance_conflict`.
188/// RecordBatch clones are cheap (column data is Arc'd).
189pub async fn merge_insert_batch_with_lance_conflict_retry(
190    backend: &dyn crate::backend::StorageBackend,
191    table_name: &str,
192    batch: arrow_array::RecordBatch,
193    on: &[&str],
194) -> anyhow::Result<()> {
195    // NOTE: `backend.merge_insert` already takes the per-table write lock
196    // internally (the same mutex `lock_table_for_write` exposes), so it is
197    // serialized against a compaction that holds that lock across its whole
198    // scan → overwrite. Do NOT take `lock_table_for_write` here as well — that
199    // would re-lock the same non-reentrant mutex and self-deadlock.
200    retry_on_lance_conflict(|| async {
201        let exists = backend.table_exists(table_name).await?;
202        if !exists {
203            anyhow::bail!(
204                "merge_insert target table '{}' does not exist (partial writes \
205                 require the row to already be present; CREATE goes through Append)",
206                table_name
207            );
208        }
209        backend
210            .merge_insert(table_name, on, vec![batch.clone()])
211            .await
212    })
213    .await
214}
215
216/// Race-safe write: creates the table if missing, otherwise appends.
217/// Each attempt re-checks `table_exists` and adjusts strategy: Append if
218/// now-exists, Create if still-missing. Retries on Lance commit conflicts
219/// via `retry_on_lance_conflict`.
220///
221/// Used by every dataset's `write_batch` helper to absorb the Lance
222/// commit-conflict-resolver behavior. RecordBatch clones are cheap
223/// (column data is Arc'd).
224pub async fn write_batch_with_lance_conflict_retry(
225    backend: &dyn crate::backend::StorageBackend,
226    table_name: &str,
227    batch: arrow_array::RecordBatch,
228) -> anyhow::Result<()> {
229    use crate::backend::types::WriteMode;
230    retry_on_lance_conflict(|| async {
231        let exists = backend.table_exists(table_name).await?;
232        if exists {
233            backend
234                .write(table_name, vec![batch.clone()], WriteMode::Append)
235                .await
236        } else {
237            backend.create_table(table_name, vec![batch.clone()]).await
238        }
239    })
240    .await
241}
242
243/// Helper to manage compaction_in_progress flag
244struct CompactionGuard {
245    status: Arc<Mutex<CompactionStatus>>,
246}
247
248impl CompactionGuard {
249    fn new(status: Arc<Mutex<CompactionStatus>>) -> Option<Self> {
250        let mut s = acquire_mutex(&status, "compaction_status").ok()?;
251        if s.compaction_in_progress {
252            return None;
253        }
254        s.compaction_in_progress = true;
255        Some(Self {
256            status: status.clone(),
257        })
258    }
259}
260
261impl Drop for CompactionGuard {
262    fn drop(&mut self) {
263        // CRITICAL: Never panic in Drop - panicking in drop() = process ABORT.
264        // See issue #18/#150. If the lock is poisoned, log and continue gracefully.
265        match uni_common::sync::acquire_mutex(&self.status, "compaction_status") {
266            Ok(mut s) => {
267                s.compaction_in_progress = false;
268                s.last_compaction = Some(std::time::SystemTime::now());
269            }
270            Err(e) => {
271                // Lock is poisoned but we're in Drop - cannot panic.
272                // Log the error and continue. System state may be inconsistent but at least
273                // we don't abort the process.
274                log::error!(
275                    "CompactionGuard drop failed to acquire poisoned lock: {}. \
276                     Compaction status may be inconsistent. Issue #18/#150",
277                    e
278                );
279            }
280        }
281    }
282}
283
284impl StorageManager {
285    /// Create a new StorageManager with a pre-configured backend.
286    pub async fn new_with_backend(
287        base_uri: &str,
288        store: Arc<dyn ObjectStore>,
289        backend: Arc<dyn StorageBackend>,
290        schema_manager: Arc<SchemaManager>,
291        config: UniConfig,
292    ) -> Result<Self> {
293        let resilient_store: Arc<dyn ObjectStore> = Arc::new(ResilientObjectStore::new(
294            store,
295            config.object_store.clone(),
296        ));
297
298        let snapshot_manager = Arc::new(SnapshotManager::new(resilient_store.clone()));
299
300        // Perform crash recovery for all known table patterns
301        Self::recover_all_staging_tables(backend.as_ref(), &schema_manager).await?;
302
303        let mut sm = Self {
304            base_uri: base_uri.to_string(),
305            store: resilient_store,
306            schema_manager,
307            snapshot_manager,
308            adjacency_manager: Arc::new(AdjacencyManager::new(config.cache_size)),
309            config,
310            compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
311            flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
312            pinned_snapshot: None,
313            pinned_version_hwm: None,
314            fork_scope: None,
315            backend,
316            vid_labels_index: Arc::new(parking_lot::RwLock::new(
317                crate::storage::vid_labels::VidLabelsIndex::new(),
318            )),
319            plugin_registry: None,
320        };
321
322        // Rebuild VidLabelsIndex from persisted vertices. A failure leaves the
323        // empty index in place; flush-time updates then repopulate it
324        // incrementally, so reads degrade rather than break.
325        if let Err(e) = sm.rebuild_vid_labels_index().await {
326            warn!(
327                "Failed to rebuild VidLabelsIndex on startup: {}. Falling back to storage queries.",
328                e
329            );
330        }
331
332        Ok(sm)
333    }
334
335    /// Create a new StorageManager with LanceDB integration.
336    #[cfg(feature = "lance-backend")]
337    pub async fn new(base_uri: &str, schema_manager: Arc<SchemaManager>) -> Result<Self> {
338        Self::new_with_config(base_uri, schema_manager, UniConfig::default()).await
339    }
340
341    /// Create a new StorageManager with custom cache size.
342    #[cfg(feature = "lance-backend")]
343    pub async fn new_with_cache(
344        base_uri: &str,
345        schema_manager: Arc<SchemaManager>,
346        adjacency_cache_size: usize,
347    ) -> Result<Self> {
348        let config = UniConfig {
349            cache_size: adjacency_cache_size,
350            ..Default::default()
351        };
352        Self::new_with_config(base_uri, schema_manager, config).await
353    }
354
355    /// Create a new StorageManager with custom configuration.
356    #[cfg(feature = "lance-backend")]
357    pub async fn new_with_config(
358        base_uri: &str,
359        schema_manager: Arc<SchemaManager>,
360        config: UniConfig,
361    ) -> Result<Self> {
362        let store = Self::build_store_from_uri(base_uri)?;
363        Self::new_with_store_and_config(base_uri, store, schema_manager, config).await
364    }
365
366    /// Create a new StorageManager using an already-constructed object store.
367    #[cfg(feature = "lance-backend")]
368    pub async fn new_with_store_and_config(
369        base_uri: &str,
370        store: Arc<dyn ObjectStore>,
371        schema_manager: Arc<SchemaManager>,
372        config: UniConfig,
373    ) -> Result<Self> {
374        Self::new_with_store_and_storage_options(base_uri, store, schema_manager, config, None)
375            .await
376    }
377
378    /// Create a new StorageManager with LanceDB storage options.
379    #[cfg(feature = "lance-backend")]
380    pub async fn new_with_store_and_storage_options(
381        base_uri: &str,
382        store: Arc<dyn ObjectStore>,
383        schema_manager: Arc<SchemaManager>,
384        config: UniConfig,
385        lancedb_storage_options: Option<HashMap<String, String>>,
386    ) -> Result<Self> {
387        let backend = Arc::new(LanceDbBackend::connect(base_uri, lancedb_storage_options).await?);
388        Self::new_with_backend(base_uri, store, backend, schema_manager, config).await
389    }
390
391    /// Recover all staging tables for known table patterns.
392    ///
393    /// This runs on startup to handle crash recovery. It checks for staging tables
394    /// for all vertex labels, adjacency tables, delta tables, and main tables.
395    async fn recover_all_staging_tables(
396        backend: &dyn StorageBackend,
397        schema_manager: &SchemaManager,
398    ) -> Result<()> {
399        let schema = schema_manager.schema();
400
401        // Recover main vertex and edge tables
402        backend
403            .recover_staging(table_names::main_vertex_table_name())
404            .await?;
405        backend
406            .recover_staging(table_names::main_edge_table_name())
407            .await?;
408
409        // Recover per-label vertex tables
410        for label in schema.labels.keys() {
411            let name = table_names::vertex_table_name(label);
412            backend.recover_staging(&name).await?;
413        }
414
415        // Recover adjacency and delta tables for each edge type and direction
416        for edge_type in schema.edge_types.keys() {
417            for direction in &["fwd", "bwd"] {
418                // Recover delta tables
419                let delta_name = table_names::delta_table_name(edge_type, direction);
420                backend.recover_staging(&delta_name).await?;
421
422                // Recover adjacency tables for each label
423                for _label in schema.labels.keys() {
424                    let adj_name = table_names::adjacency_table_name(edge_type, direction);
425                    backend.recover_staging(&adj_name).await?;
426                }
427            }
428        }
429
430        Ok(())
431    }
432
433    #[cfg(feature = "lance-backend")]
434    fn build_store_from_uri(base_uri: &str) -> Result<Arc<dyn ObjectStore>> {
435        if base_uri.contains("://") {
436            let parsed = url::Url::parse(base_uri).map_err(|e| anyhow!("Invalid base URI: {e}"))?;
437            let (store, _path) = object_store::parse_url(&parsed)
438                .map_err(|e| anyhow!("Failed to parse object store URL: {e}"))?;
439            Ok(Arc::from(store))
440        } else {
441            // If local path, ensure it exists.
442            std::fs::create_dir_all(base_uri)?;
443            Ok(Arc::new(LocalFileSystem::new_with_prefix(base_uri)?))
444        }
445    }
446
447    /// Filesystem root backing this manager's object store, when the store
448    /// is a local filesystem (the non-`://` branch of
449    /// `build_store_from_uri`). Used to fsync WAL segments after PUT —
450    /// `object_store::LocalFileSystem` does not fsync on its own. `None`
451    /// for remote/URL-based stores.
452    pub fn local_fs_root(&self) -> Option<std::path::PathBuf> {
453        if self.base_uri.contains("://") {
454            None
455        } else {
456            Some(std::path::PathBuf::from(&self.base_uri))
457        }
458    }
459
460    pub fn pinned(&self, snapshot: SnapshotManifest) -> Self {
461        // Phase 4a: pinning a forked session is now supported. The
462        // resulting StorageManager keeps `fork_scope` so reads continue
463        // to route through the fork's Lance branches via `base_paths`,
464        // and adds `pinned_snapshot` so writers / writers' read views
465        // resolve at the snapshot's HWM. Writes are gated separately by
466        // the session-level `is_pinned` check (`Session::tx` rejects
467        // them via `UniError::ReadOnly`).
468        Self {
469            base_uri: self.base_uri.clone(),
470            store: self.store.clone(),
471            schema_manager: self.schema_manager.clone(),
472            snapshot_manager: self.snapshot_manager.clone(),
473            // Separate AdjacencyManager for snapshot isolation (Issue #73):
474            // warm() will load only edges visible at the snapshot's HWM.
475            // This prevents live DB's CSR (with all edges) from leaking into snapshots.
476            adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
477            config: self.config.clone(),
478            compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
479            flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
480            pinned_snapshot: Some(snapshot),
481            pinned_version_hwm: None,
482            fork_scope: self.fork_scope.clone(),
483            backend: self.backend.clone(),
484            // Deep-copy, not Arc-clone: a fork/pin must get its OWN label index
485            // so its flushes/relabels don't mutate the parent's (review H1/L2),
486            // mirroring the fresh `adjacency_manager` above. `VidLabelsIndex`
487            // derives `Clone`; the snapshot is taken after flush-before-branch so
488            // inherited labels (#99) are preserved.
489            vid_labels_index: Arc::new(parking_lot::RwLock::new(
490                self.vid_labels_index.read().clone(),
491            )),
492            plugin_registry: self.plugin_registry.clone(),
493        }
494    }
495
496    /// Construct a clone of this `StorageManager` pinned to a row-version
497    /// high-water mark (C2: transaction-level L1 pinning).
498    ///
499    /// Unlike [`Self::pinned`], this needs no `SnapshotManifest`: scans
500    /// filter to `_version <= hwm` via [`Self::version_high_water_mark`].
501    /// A read-write transaction builds one of these at begin with
502    /// `SnapshotView.started_at_version`, so an L0→L1 flush completing
503    /// mid-transaction cannot leak post-snapshot rows into its L1 scans
504    /// (the L0 tier is pinned separately by the `SnapshotView`).
505    ///
506    /// Unlike [`Self::pinned`], the live `AdjacencyManager` is SHARED, not
507    /// fresh: commits replay their edges into the live manager's overlay,
508    /// which is the traversal path's only source for L0-resident edges — a
509    /// fresh manager would make every unflushed edge invisible to the
510    /// transaction. The cost is that the edge tier is not version-pinned
511    /// (post-snapshot edges remain visible to traversals, exactly as before
512    /// C2); edge reads are recorded in the OCC read-set, so a conflicting
513    /// read-modify-write still aborts at commit.
514    pub fn pinned_at_version(&self, hwm: u64) -> Self {
515        Self {
516            base_uri: self.base_uri.clone(),
517            store: self.store.clone(),
518            schema_manager: self.schema_manager.clone(),
519            snapshot_manager: self.snapshot_manager.clone(),
520            adjacency_manager: self.adjacency_manager.clone(),
521            config: self.config.clone(),
522            compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
523            flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
524            pinned_snapshot: None,
525            pinned_version_hwm: Some(hwm),
526            fork_scope: self.fork_scope.clone(),
527            backend: self.backend.clone(),
528            // Deep-copy, not Arc-clone: a fork/pin must get its OWN label index
529            // so its flushes/relabels don't mutate the parent's (review H1/L2),
530            // mirroring the fresh `adjacency_manager` above. `VidLabelsIndex`
531            // derives `Clone`; the snapshot is taken after flush-before-branch so
532            // inherited labels (#99) are preserved.
533            vid_labels_index: Arc::new(parking_lot::RwLock::new(
534                self.vid_labels_index.read().clone(),
535            )),
536            plugin_registry: self.plugin_registry.clone(),
537        }
538    }
539
540    /// Construct a fork-scoped clone of this `StorageManager`.
541    ///
542    /// All reads through dataset factories *and* through `backend()`
543    /// on the returned manager route through the fork's Lance branches
544    /// via `base_paths`. The `AdjacencyManager` is fresh (per Issue
545    /// #73 reasoning — same as `pinned`) to prevent primary's CSR from
546    /// leaking into the fork. `fork_scope` and `pinned_snapshot` are
547    /// mutually exclusive.
548    ///
549    /// The backend is wrapped in [`crate::backend::branched::BranchedBackend`]
550    /// so that every `ScanRequest` constructed *anywhere* (PropertyManager,
551    /// MainVertexDataset static methods, etc.) automatically picks up
552    /// the fork's branch for tables the fork has branched. Untracked
553    /// tables fall back to primary, matching Phase 1 read semantics.
554    pub fn at_fork(&self, scope: Arc<crate::fork::ForkScope>) -> Self {
555        self.at_fork_with_schema(scope, self.schema_manager.clone())
556    }
557
558    /// Variant of [`Self::at_fork`] that uses an explicit
559    /// `merged_schema` for the fork's storage rather than primary's
560    /// schema_manager. Used by `UniInner::at_fork` so that the
561    /// fork-side strict-schema checks (in `uni-query` / `uni-store`'s
562    /// writer) see fork-local labels and edge types added through
563    /// `Session::fork_schema()`. Without this, those checks would
564    /// route through primary's schema and reject fork-local labels.
565    pub fn at_fork_with_schema(
566        &self,
567        scope: Arc<crate::fork::ForkScope>,
568        merged_schema: Arc<SchemaManager>,
569    ) -> Self {
570        debug_assert!(
571            self.pinned_snapshot.is_none(),
572            "forking a pinned StorageManager is unsupported in Phase 1"
573        );
574        let branched_backend: Arc<dyn StorageBackend> = Arc::new(
575            crate::backend::branched::BranchedBackend::new(self.backend.clone(), scope.clone()),
576        );
577        // Fork-scoped snapshot manager: a fork's flush publishes its manifest +
578        // `latest` pointer under `catalog/forks/{fork_id}/`, never the primary's
579        // global `catalog/latest` (review C1). Uses the raw object store, since
580        // catalog metadata is not branched.
581        let snapshot_manager = Arc::new(SnapshotManager::new_for_fork(
582            self.store.clone(),
583            scope.fork_id(),
584        ));
585        Self {
586            base_uri: self.base_uri.clone(),
587            store: self.store.clone(),
588            schema_manager: merged_schema,
589            snapshot_manager,
590            adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
591            config: self.config.clone(),
592            compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
593            flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
594            pinned_snapshot: None,
595            pinned_version_hwm: None,
596            fork_scope: Some(scope),
597            backend: branched_backend,
598            // Deep-copy, not Arc-clone: a fork/pin must get its OWN label index
599            // so its flushes/relabels don't mutate the parent's (review H1/L2),
600            // mirroring the fresh `adjacency_manager` above. `VidLabelsIndex`
601            // derives `Clone`; the snapshot is taken after flush-before-branch so
602            // inherited labels (#99) are preserved.
603            vid_labels_index: Arc::new(parking_lot::RwLock::new(
604                self.vid_labels_index.read().clone(),
605            )),
606            plugin_registry: self.plugin_registry.clone(),
607        }
608    }
609
610    /// Borrow the plugin registry used for registry-dispatched CRDT merges.
611    ///
612    /// Returns `None` when no registry has been installed, in which case the
613    /// durable merge paths fall back to native [`uni_crdt::Crdt::try_merge`].
614    pub fn plugin_registry(&self) -> Option<&Arc<uni_plugin::PluginRegistry>> {
615        self.plugin_registry.as_ref()
616    }
617
618    /// Install the plugin registry used for registry-dispatched CRDT merges.
619    ///
620    /// Called once at DB construction (before the manager is shared) so the
621    /// same registry that backs `PropertyManager` also governs the compaction
622    /// and L0-flush durable merge paths.
623    pub fn set_plugin_registry(&mut self, registry: Arc<uni_plugin::PluginRegistry>) {
624        self.plugin_registry = Some(registry);
625    }
626
627    /// Borrow the active fork scope, if any.
628    pub fn fork_scope(&self) -> Option<&Arc<crate::fork::ForkScope>> {
629        self.fork_scope.as_ref()
630    }
631
632    /// Phase 5a: query whether a fork-local index of `kind` exists
633    /// for the `(label, column)` pair on the active fork scope.
634    /// Returns `false` outside a fork or when no fork-local build of
635    /// that kind has completed for the pair.
636    ///
637    /// The planner consults this to decide whether to emit a specific
638    /// `FusedIndexScan` (returns `true`) or fall back to the inherited
639    /// primary index via `base_paths` (returns `false`). A column can
640    /// carry several kinds at once, so callers ask for the exact kind
641    /// they intend to fuse. The lookup is a `DashMap::get` on
642    /// `ForkScope` — O(1) and safe to call per query without caching
643    /// above this layer.
644    #[must_use]
645    pub fn has_fork_index(
646        &self,
647        label: &str,
648        column: &str,
649        kind: crate::fork::ForkLocalIndexKind,
650    ) -> bool {
651        self.fork_scope
652            .as_ref()
653            .is_some_and(|s| s.has_fork_local_index(label, column, kind))
654    }
655
656    /// Base URI for this storage manager (the directory or remote
657    /// prefix under which dataset directories live).
658    pub fn base_uri(&self) -> &str {
659        &self.base_uri
660    }
661
662    pub fn get_edge_version_by_id(&self, edge_type_id: u32) -> Option<u64> {
663        let schema = self.schema_manager.schema();
664        let name = schema.edge_type_name_by_id(edge_type_id)?;
665        self.pinned_snapshot
666            .as_ref()
667            .and_then(|s| s.edges.get(name).map(|es| es.lance_version))
668            // The flush path stamps `lance_version: 0` ("LanceDB tables don't
669            // expose Lance version directly") — 0 is a stub sentinel, not a
670            // real dataset version. Returning it would route adjacency reads
671            // through `checkout_version(0)` (the empty initial version).
672            .filter(|v| *v != 0)
673    }
674
675    /// Returns the version high-water mark from the pinned snapshot or the
676    /// transaction-level version pin, if present.
677    ///
678    /// Used by the SCAN tier (vertex tables, property reads) to filter data
679    /// by version: when set, only rows with
680    /// `_version <= version_high_water_mark` are visible. The edge/adjacency
681    /// path must use [`Self::snapshot_version_hwm`] instead.
682    pub fn version_high_water_mark(&self) -> Option<u64> {
683        self.pinned_snapshot
684            .as_ref()
685            .map(|s| s.version_high_water_mark)
686            .or(self.pinned_version_hwm)
687    }
688
689    /// Version high-water mark from a manifest-pinned (time-travel) snapshot
690    /// ONLY — never from a transaction-level version pin.
691    ///
692    /// The edge/adjacency read path switches to version-filtered CSR reads
693    /// and skips the L0 overlays when a hwm is present. That is correct for
694    /// time-travel (a snapshot is flushed state, with its own fresh
695    /// `AdjacencyManager`), but a transaction pin shares the LIVE adjacency
696    /// manager and needs live CSR + L0 overlays + its tx-L0 — filtering
697    /// there would hide unflushed edges and poison the shared warm cache.
698    /// The edge tier is deliberately not version-pinned for transactions
699    /// (see [`Self::pinned_at_version`]).
700    pub fn snapshot_version_hwm(&self) -> Option<u64> {
701        self.pinned_snapshot
702            .as_ref()
703            .map(|s| s.version_high_water_mark)
704    }
705
706    /// Apply version filtering to a base filter expression.
707    ///
708    /// If a snapshot is pinned, wraps `base_filter` with an additional
709    /// `_version <= hwm` clause. Otherwise returns `base_filter` unchanged.
710    pub fn apply_version_filter(&self, base_filter: String) -> String {
711        if let Some(hwm) = self.version_high_water_mark() {
712            format!("({}) AND (_version <= {})", base_filter, hwm)
713        } else {
714            base_filter
715        }
716    }
717
718    /// Build a filter expression that excludes soft-deleted rows and optionally
719    /// includes a user-provided filter.
720    fn build_active_filter(user_filter: Option<&str>) -> String {
721        match user_filter {
722            Some(expr) => format!("({}) AND (_deleted = false)", expr),
723            None => "_deleted = false".to_string(),
724        }
725    }
726
727    pub fn store(&self) -> Arc<dyn ObjectStore> {
728        self.store.clone()
729    }
730
731    /// Get current compaction status.
732    ///
733    /// # Errors
734    ///
735    /// Returns error if the compaction status lock is poisoned (see issue #18/#150).
736    pub fn compaction_status(
737        &self,
738    ) -> Result<CompactionStatus, uni_common::sync::LockPoisonedError> {
739        let guard = uni_common::sync::acquire_mutex(&self.compaction_status, "compaction_status")?;
740        Ok(guard.clone())
741    }
742
743    pub async fn compact(&self) -> Result<CompactionStats> {
744        // Backend handles compaction internally via optimize_table()
745        let start = std::time::Instant::now();
746        let schema = self.schema_manager.schema();
747        let mut files_compacted = 0;
748
749        for label in schema.labels.keys() {
750            let name = table_names::vertex_table_name(label);
751            if self.backend.table_exists(&name).await? {
752                self.backend.optimize_table(&name).await?;
753                files_compacted += 1;
754                self.backend.invalidate_cache(&name);
755            }
756        }
757
758        Ok(CompactionStats {
759            files_compacted,
760            bytes_before: 0,
761            bytes_after: 0,
762            duration: start.elapsed(),
763            crdt_merges: 0,
764        })
765    }
766
767    pub async fn compact_label(&self, label: &str) -> Result<CompactionStats> {
768        let _guard = CompactionGuard::new(self.compaction_status.clone())
769            .ok_or_else(|| anyhow!("Compaction already in progress"))?;
770
771        let start = std::time::Instant::now();
772        let name = table_names::vertex_table_name(label);
773
774        if self.backend.table_exists(&name).await? {
775            self.backend.optimize_table(&name).await?;
776            self.backend.invalidate_cache(&name);
777        }
778
779        Ok(CompactionStats {
780            files_compacted: 1,
781            bytes_before: 0,
782            bytes_after: 0,
783            duration: start.elapsed(),
784            crdt_merges: 0,
785        })
786    }
787
788    pub async fn compact_edge_type(&self, edge_type: &str) -> Result<CompactionStats> {
789        let _guard = CompactionGuard::new(self.compaction_status.clone())
790            .ok_or_else(|| anyhow!("Compaction already in progress"))?;
791
792        let start = std::time::Instant::now();
793        let mut files_compacted = 0;
794
795        for dir in ["fwd", "bwd"] {
796            let name = table_names::delta_table_name(edge_type, dir);
797            if self.backend.table_exists(&name).await? {
798                self.backend.optimize_table(&name).await?;
799                files_compacted += 1;
800            }
801        }
802
803        Ok(CompactionStats {
804            files_compacted,
805            bytes_before: 0,
806            bytes_after: 0,
807            duration: start.elapsed(),
808            crdt_merges: 0,
809        })
810    }
811
812    pub async fn wait_for_compaction(&self) -> Result<()> {
813        loop {
814            let in_progress = {
815                acquire_mutex(&self.compaction_status, "compaction_status")?.compaction_in_progress
816            };
817            if !in_progress {
818                return Ok(());
819            }
820            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
821        }
822    }
823
824    pub fn start_background_compaction(
825        self: Arc<Self>,
826        mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
827    ) -> tokio::task::JoinHandle<()> {
828        if !self.config.compaction.enabled {
829            return tokio::spawn(async {});
830        }
831
832        tokio::spawn(async move {
833            // Use interval_at to delay the first tick. tokio::time::interval fires
834            // immediately on the first tick, which can race with queries that run
835            // right after database open. Delaying by the check_interval gives
836            // initial queries time to complete before compaction modifies tables
837            // (optimize(All) can GC index files that concurrent queries depend on).
838            let start = tokio::time::Instant::now() + self.config.compaction.check_interval;
839            let mut interval =
840                tokio::time::interval_at(start, self.config.compaction.check_interval);
841
842            loop {
843                tokio::select! {
844                    _ = interval.tick() => {
845                        if let Err(e) = self.update_compaction_status().await {
846                            log::error!("Failed to update compaction status: {}", e);
847                            continue;
848                        }
849
850                        if let Some(task) = self.pick_compaction_task() {
851                            log::info!("Triggering background compaction: {:?}", task);
852                            if let Err(e) = Self::execute_compaction(Arc::clone(&self), task).await {
853                                log::error!("Compaction failed: {}", e);
854                            }
855                        }
856                    }
857                    _ = shutdown_rx.recv() => {
858                        log::info!("Background compaction shutting down");
859                        let _ = self.wait_for_compaction().await;
860                        break;
861                    }
862                }
863            }
864        })
865    }
866
867    async fn update_compaction_status(&self) -> Result<()> {
868        let schema = self.schema_manager.schema();
869        let backend = self.backend.as_ref();
870        let mut total_rows: usize = 0;
871        let mut oldest_ts: Option<i64> = None;
872
873        for name in schema.edge_types.keys() {
874            for dir in ["fwd", "bwd"] {
875                let tbl_name = table_names::delta_table_name(name, dir);
876                if !backend.table_exists(&tbl_name).await? {
877                    continue;
878                }
879                let row_count = backend.count_rows(&tbl_name, None).await.unwrap_or(0);
880                if row_count == 0 {
881                    continue;
882                }
883                total_rows += row_count;
884
885                // Query oldest _created_at for age tracking
886                let request =
887                    ScanRequest::all(&tbl_name).with_columns(vec!["_created_at".to_string()]);
888                let Ok(batches) = backend.scan(request).await else {
889                    continue;
890                };
891                for batch in batches {
892                    let Some(col) = batch
893                        .column_by_name("_created_at")
894                        .and_then(|c| c.as_any().downcast_ref::<TimestampNanosecondArray>())
895                    else {
896                        continue;
897                    };
898                    for i in 0..col.len() {
899                        if !col.is_null(i) {
900                            let ts = col.value(i);
901                            oldest_ts = Some(oldest_ts.map_or(ts, |prev| prev.min(ts)));
902                        }
903                    }
904                }
905            }
906        }
907
908        let oldest_l1_age = oldest_ts
909            .and_then(|ts| {
910                let created = UNIX_EPOCH + Duration::from_nanos(ts as u64);
911                SystemTime::now().duration_since(created).ok()
912            })
913            .unwrap_or(Duration::ZERO);
914
915        let mut status = acquire_mutex(&self.compaction_status, "compaction_status")?;
916        // Note: l1_runs is managed by flush_to_l1 (increment) and execute_compaction
917        // (reset). It counts flush generations, not delta table count.
918        status.l1_size_bytes = (total_rows * ENTRY_SIZE_ESTIMATE) as u64;
919        status.oldest_l1_age = oldest_l1_age;
920        Ok(())
921    }
922
923    fn pick_compaction_task(&self) -> Option<CompactionTask> {
924        let status = acquire_mutex(&self.compaction_status, "compaction_status").ok()?;
925
926        if status.l1_runs >= self.config.compaction.max_l1_runs {
927            return Some(CompactionTask::ByRunCount);
928        }
929        if status.l1_size_bytes >= self.config.compaction.max_l1_size_bytes {
930            return Some(CompactionTask::BySize);
931        }
932        if status.oldest_l1_age >= self.config.compaction.max_l1_age
933            && status.oldest_l1_age > Duration::ZERO
934        {
935            return Some(CompactionTask::ByAge);
936        }
937
938        None
939    }
940
941    /// Optimize a table via the backend, returning `true` on success.
942    async fn try_optimize_table(backend: &dyn StorageBackend, table_name: &str) -> bool {
943        if let Err(e) = backend.optimize_table(table_name).await {
944            log::warn!("Failed to optimize table {}: {}", table_name, e);
945            return false;
946        }
947        true
948    }
949
950    /// Trigger L1 compaction asynchronously without blocking the caller.
951    /// Safe to call frequently — CompactionGuard prevents concurrent runs.
952    pub fn trigger_async_compaction(self: &Arc<Self>) {
953        let this = Arc::clone(self);
954        tokio::spawn(async move {
955            if let Err(e) = Self::execute_compaction(this, CompactionTask::ByRunCount).await {
956                // "Compaction already in progress" is expected when called frequently
957                log::debug!("Post-flush compaction skipped: {}", e);
958            }
959        });
960    }
961
962    pub(crate) async fn execute_compaction(
963        this: Arc<Self>,
964        _task: CompactionTask,
965    ) -> Result<CompactionStats> {
966        let start = std::time::Instant::now();
967        let _guard = CompactionGuard::new(this.compaction_status.clone())
968            .ok_or_else(|| anyhow!("Compaction already in progress"))?;
969
970        let schema = this.schema_manager.schema();
971        let mut files_compacted = 0;
972
973        // ── Tier 2: Semantic compaction ──
974        // Dedup vertices, merge CRDTs, consolidate L1→L2 deltas, clean tombstones
975        let compactor = Compactor::new(Arc::clone(&this));
976        let compaction_results = compactor.compact_all().await.unwrap_or_else(|e| {
977            log::error!(
978                "Semantic compaction failed (continuing with backend optimize): {}",
979                e
980            );
981            Vec::new()
982        });
983
984        // Re-warm adjacency CSR after semantic compaction
985        let am = this.adjacency_manager();
986        for info in &compaction_results {
987            let direction = match info.direction.as_str() {
988                "fwd" => Direction::Outgoing,
989                "bwd" => Direction::Incoming,
990                _ => continue,
991            };
992            if let Some(etid) = schema.edge_type_id_unified_case_insensitive(&info.edge_type)
993                && let Err(e) = am.warm(&this, etid, direction, None).await
994            {
995                log::warn!(
996                    "Failed to re-warm adjacency for {}/{}: {}",
997                    info.edge_type,
998                    info.direction,
999                    e
1000                );
1001            }
1002        }
1003
1004        // ── Tier 3: Backend optimize ──
1005        let backend = this.backend.as_ref();
1006
1007        // Optimize edge delta and adjacency tables
1008        for name in schema.edge_types.keys() {
1009            for dir in ["fwd", "bwd"] {
1010                let delta = table_names::delta_table_name(name, dir);
1011                if Self::try_optimize_table(backend, &delta).await {
1012                    files_compacted += 1;
1013                }
1014                let adj = table_names::adjacency_table_name(name, dir);
1015                if Self::try_optimize_table(backend, &adj).await {
1016                    files_compacted += 1;
1017                }
1018            }
1019        }
1020
1021        // Optimize vertex tables
1022        for label in schema.labels.keys() {
1023            let tbl = table_names::vertex_table_name(label);
1024            if Self::try_optimize_table(backend, &tbl).await {
1025                files_compacted += 1;
1026                backend.invalidate_cache(&tbl);
1027            }
1028        }
1029
1030        // Optimize main vertex and edge tables
1031        for tbl in [
1032            table_names::main_vertex_table_name(),
1033            table_names::main_edge_table_name(),
1034        ] {
1035            if Self::try_optimize_table(backend, tbl).await {
1036                files_compacted += 1;
1037            }
1038        }
1039
1040        {
1041            let mut status = acquire_mutex(&this.compaction_status, "compaction_status")?;
1042            status.total_compactions += 1;
1043            status.l1_runs = 0; // Reset flush generation counter
1044        }
1045
1046        Ok(CompactionStats {
1047            files_compacted,
1048            bytes_before: 0,
1049            bytes_after: 0,
1050            duration: start.elapsed(),
1051            crdt_merges: 0,
1052        })
1053    }
1054
1055    /// Open a LanceDB table for a label.
1056    ///
1057    /// Invalidate cached table state (call after writes).
1058    pub fn invalidate_table_cache(&self, label: &str) {
1059        let name = table_names::vertex_table_name(label);
1060        self.backend.invalidate_cache(&name);
1061    }
1062
1063    pub fn base_path(&self) -> &str {
1064        &self.base_uri
1065    }
1066
1067    pub fn schema_manager(&self) -> &SchemaManager {
1068        &self.schema_manager
1069    }
1070
1071    pub fn schema_manager_arc(&self) -> Arc<SchemaManager> {
1072        self.schema_manager.clone()
1073    }
1074
1075    /// Returns the backing `Arc<SchemaManager>` by reference.
1076    ///
1077    /// Unlike [`Self::schema_manager`] (which derefs to `&SchemaManager`),
1078    /// this preserves the `Arc`'s pointer identity. A pinned transaction and
1079    /// the live session clone the *same* `schema_manager` `Arc`, while forks
1080    /// hold a distinct one — so this is the correct registry key for the
1081    /// projection store (see `uni-query`'s `projection_store::for_storage`).
1082    #[must_use]
1083    pub fn schema_manager_arc_ref(&self) -> &Arc<SchemaManager> {
1084        &self.schema_manager
1085    }
1086
1087    /// Get the adjacency manager for the dual-CSR architecture.
1088    pub fn adjacency_manager(&self) -> Arc<AdjacencyManager> {
1089        Arc::clone(&self.adjacency_manager)
1090    }
1091
1092    /// Warm the adjacency manager for a specific edge type and direction.
1093    ///
1094    /// Builds the Main CSR from L2 adjacency + L1 delta data in storage.
1095    /// Called lazily on first access per edge type or at startup.
1096    pub async fn warm_adjacency(
1097        &self,
1098        edge_type_id: u32,
1099        direction: crate::storage::direction::Direction,
1100        version: Option<u64>,
1101    ) -> anyhow::Result<()> {
1102        self.adjacency_manager
1103            .warm(self, edge_type_id, direction, version)
1104            .await
1105    }
1106
1107    /// Coalesced warm_adjacency() to prevent cache stampede (Issue #13).
1108    ///
1109    /// Uses double-checked locking to ensure only one concurrent warm() per
1110    /// (edge_type, direction) key. Subsequent callers wait for the first to complete.
1111    pub async fn warm_adjacency_coalesced(
1112        &self,
1113        edge_type_id: u32,
1114        direction: crate::storage::direction::Direction,
1115        version: Option<u64>,
1116    ) -> anyhow::Result<()> {
1117        self.adjacency_manager
1118            .warm_coalesced(self, edge_type_id, direction, version)
1119            .await
1120    }
1121
1122    /// Check whether the adjacency manager has a CSR for the given edge type and direction.
1123    pub fn has_adjacency_csr(
1124        &self,
1125        edge_type_id: u32,
1126        direction: crate::storage::direction::Direction,
1127    ) -> bool {
1128        self.adjacency_manager.has_csr(edge_type_id, direction)
1129    }
1130
1131    /// Get neighbors at a specific version for snapshot queries.
1132    pub fn get_neighbors_at_version(
1133        &self,
1134        vid: uni_common::core::id::Vid,
1135        edge_type: u32,
1136        direction: crate::storage::direction::Direction,
1137        version: u64,
1138    ) -> Vec<(uni_common::core::id::Vid, uni_common::core::id::Eid)> {
1139        self.adjacency_manager
1140            .get_neighbors_at_version(vid, edge_type, direction, version)
1141    }
1142
1143    /// Get the storage backend.
1144    pub fn backend(&self) -> &dyn StorageBackend {
1145        self.backend.as_ref()
1146    }
1147
1148    /// Get the storage backend as an Arc.
1149    pub fn backend_arc(&self) -> Arc<dyn StorageBackend> {
1150        self.backend.clone()
1151    }
1152
1153    /// Rebuild the VidLabelsIndex from the main vertex table.
1154    ///
1155    /// Always called on startup. On a fresh database (no vertex table yet) the
1156    /// index is left empty and filled incrementally by flush-time updates.
1157    async fn rebuild_vid_labels_index(&mut self) -> Result<()> {
1158        use crate::storage::vid_labels::VidLabelsIndex;
1159
1160        let backend = self.backend.as_ref();
1161        let vtable = table_names::main_vertex_table_name();
1162
1163        // Check if the table exists (fresh database)
1164        if !backend.table_exists(vtable).await? {
1165            self.vid_labels_index = Arc::new(parking_lot::RwLock::new(VidLabelsIndex::new()));
1166            return Ok(());
1167        }
1168
1169        // Scan all non-deleted vertices and collect (VID, labels)
1170        let request = ScanRequest::all(vtable)
1171            .with_filter("_deleted = false")
1172            .with_limit(100_000);
1173        let batches = backend
1174            .scan(request)
1175            .await
1176            .map_err(|e| anyhow!("Failed to query main vertex table: {}", e))?;
1177
1178        let mut index = VidLabelsIndex::new();
1179        for batch in batches {
1180            let vid_col = batch
1181                .column_by_name("_vid")
1182                .ok_or_else(|| anyhow!("Missing _vid column"))?
1183                .as_any()
1184                .downcast_ref::<UInt64Array>()
1185                .ok_or_else(|| anyhow!("Invalid _vid column type"))?;
1186
1187            let labels_col = batch
1188                .column_by_name("labels")
1189                .ok_or_else(|| anyhow!("Missing labels column"))?
1190                .as_any()
1191                .downcast_ref::<arrow_array::ListArray>()
1192                .ok_or_else(|| anyhow!("Invalid labels column type"))?;
1193
1194            for row_idx in 0..batch.num_rows() {
1195                let vid = Vid::from(vid_col.value(row_idx));
1196                let labels_array = labels_col.value(row_idx);
1197                let labels_str_array = labels_array
1198                    .as_any()
1199                    .downcast_ref::<arrow_array::StringArray>()
1200                    .ok_or_else(|| anyhow!("Invalid labels array element type"))?;
1201
1202                let labels: Vec<String> = (0..labels_str_array.len())
1203                    .map(|i| labels_str_array.value(i).to_string())
1204                    .collect();
1205
1206                index.insert(vid, labels);
1207            }
1208        }
1209
1210        self.vid_labels_index = Arc::new(parking_lot::RwLock::new(index));
1211        Ok(())
1212    }
1213
1214    /// Get labels for a VID from the in-memory index.
1215    ///
1216    /// Returns `None` only when the VID is absent from the index (e.g. it was
1217    /// never persisted, or has been deleted).
1218    pub fn get_labels_from_index(&self, vid: Vid) -> Option<Vec<String>> {
1219        let index = self.vid_labels_index.read();
1220        index.get_labels(vid).map(|labels| labels.to_vec())
1221    }
1222
1223    /// Update the VID-to-labels mapping in the index.
1224    pub fn update_vid_labels_index(&self, vid: Vid, labels: Vec<String>) {
1225        let mut index = self.vid_labels_index.write();
1226        index.insert(vid, labels);
1227    }
1228
1229    /// Remove a VID from the labels index.
1230    pub fn remove_from_vid_labels_index(&self, vid: Vid) {
1231        let mut index = self.vid_labels_index.write();
1232        index.remove_vid(vid);
1233    }
1234
1235    pub async fn load_subgraph_cached(
1236        &self,
1237        start_vids: &[Vid],
1238        edge_types: &[u32],
1239        max_hops: usize,
1240        direction: GraphDirection,
1241        _l0: Option<Arc<RwLock<L0Buffer>>>,
1242    ) -> Result<WorkingGraph> {
1243        let mut graph = WorkingGraph::new();
1244
1245        let dir = match direction {
1246            GraphDirection::Outgoing => crate::storage::direction::Direction::Outgoing,
1247            GraphDirection::Incoming => crate::storage::direction::Direction::Incoming,
1248        };
1249
1250        let neighbor_is_dst = matches!(direction, GraphDirection::Outgoing);
1251
1252        // Initialize frontier
1253        let mut frontier: Vec<Vid> = start_vids.to_vec();
1254        let mut visited: HashSet<Vid> = HashSet::new();
1255
1256        // Initialize start vids
1257        for &vid in start_vids {
1258            graph.add_vertex(vid);
1259        }
1260
1261        for _hop in 0..max_hops {
1262            let mut next_frontier = HashSet::new();
1263
1264            for &vid in &frontier {
1265                if visited.contains(&vid) {
1266                    continue;
1267                }
1268                visited.insert(vid);
1269                graph.add_vertex(vid);
1270
1271                for &etype_id in edge_types {
1272                    // Warm adjacency with coalescing to prevent cache stampede (Issue #13).
1273                    // Manifest pin only: a tx version pin shares the LIVE adjacency
1274                    // manager — warming it filtered would poison the shared cache.
1275                    let edge_ver = self.snapshot_version_hwm();
1276                    self.adjacency_manager
1277                        .warm_coalesced(self, etype_id, dir, edge_ver)
1278                        .await?;
1279
1280                    // Get neighbors from AdjacencyManager (Main CSR + overlay)
1281                    let edges = self.adjacency_manager.get_neighbors(vid, etype_id, dir);
1282
1283                    for (neighbor_vid, eid) in edges {
1284                        graph.add_vertex(neighbor_vid);
1285                        if !visited.contains(&neighbor_vid) {
1286                            next_frontier.insert(neighbor_vid);
1287                        }
1288
1289                        if neighbor_is_dst {
1290                            graph.add_edge(vid, neighbor_vid, eid, etype_id);
1291                        } else {
1292                            graph.add_edge(neighbor_vid, vid, eid, etype_id);
1293                        }
1294                    }
1295                }
1296            }
1297            frontier = next_frontier.into_iter().collect();
1298
1299            // Early termination: if frontier is empty, no more vertices to explore
1300            if frontier.is_empty() {
1301                break;
1302            }
1303        }
1304
1305        Ok(graph)
1306    }
1307
1308    pub fn snapshot_manager(&self) -> &SnapshotManager {
1309        &self.snapshot_manager
1310    }
1311
1312    pub fn index_manager(&self) -> IndexManager {
1313        IndexManager::new(&self.base_uri, self.schema_manager.clone())
1314            .with_backend(self.backend_arc())
1315    }
1316
1317    // ========================================================================
1318    // Domain-level scan methods — encapsulate LanceDB queries for consumers
1319    // ========================================================================
1320
1321    /// Scan a per-label vertex table. Returns `None` if the table doesn't exist.
1322    ///
1323    /// Internally opens the table, filters requested columns to those that
1324    /// physically exist, and applies the version HWM filter for snapshot isolation.
1325    pub async fn scan_vertex_table(
1326        &self,
1327        label: &str,
1328        columns: &[&str],
1329        additional_filter: Option<&str>,
1330    ) -> Result<Option<arrow_array::RecordBatch>> {
1331        let backend = self.backend();
1332        let table_name = table_names::vertex_table_name(label);
1333
1334        if !backend.table_exists(&table_name).await? {
1335            return Ok(None);
1336        }
1337
1338        // Filter columns to those that exist in the table
1339        let actual_columns =
1340            if let Some(table_schema) = backend.get_table_schema(&table_name).await? {
1341                let table_field_names: HashSet<&str> = table_schema
1342                    .fields()
1343                    .iter()
1344                    .map(|f| f.name().as_str())
1345                    .collect();
1346                columns
1347                    .iter()
1348                    .copied()
1349                    .filter(|c| table_field_names.contains(c))
1350                    .map(|s| s.to_string())
1351                    .collect::<Vec<_>>()
1352            } else {
1353                return Ok(None);
1354            };
1355
1356        // Build filter with version HWM + optional additional filter
1357        let filter = match (self.version_high_water_mark(), additional_filter) {
1358            (Some(hwm), Some(f)) => Some(format!("_version <= {} AND ({})", hwm, f)),
1359            (Some(hwm), None) => Some(format!("_version <= {}", hwm)),
1360            (None, Some(f)) => Some(f.to_string()),
1361            (None, None) => None,
1362        };
1363
1364        let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1365        if let Some(f) = filter {
1366            request = request.with_filter(f);
1367        }
1368
1369        // Fail closed: a scan error (transient I/O, an unparsable filter, a
1370        // corrupt fragment) must propagate, never be silently mapped to
1371        // `Ok(None)`. Callers treat `Ok(None)` as "no rows" — e.g. the MERGE
1372        // fast path would create a duplicate node on a transient failure (review
1373        // bug #3a) — so an error here must surface as an error. A genuinely-
1374        // absent table is already handled above.
1375        let batches = backend.scan(request).await?;
1376        if batches.is_empty() {
1377            Ok(None)
1378        } else {
1379            Ok(Some(arrow::compute::concat_batches(
1380                &batches[0].schema(),
1381                &batches,
1382            )?))
1383        }
1384    }
1385
1386    /// Scan a delta table for an edge type + direction.
1387    /// Returns `None` if the table doesn't exist.
1388    pub async fn scan_delta_table(
1389        &self,
1390        edge_type: &str,
1391        direction: &str,
1392        columns: &[&str],
1393        additional_filter: Option<&str>,
1394    ) -> Result<Option<arrow_array::RecordBatch>> {
1395        // Edge path: manifest pin only. A transaction version pin must NOT
1396        // version-filter edge reads — the edge tier is not version-pinned
1397        // (the live AdjacencyManager + tx-L0 overlay carry unflushed and
1398        // in-transaction edges), so filtering here would hide a relationship
1399        // the same transaction just created (MERGE read-your-writes).
1400        let edge_hwm = self.snapshot_version_hwm();
1401        let backend = self.backend();
1402        let table_name = table_names::delta_table_name(edge_type, direction);
1403
1404        if !backend.table_exists(&table_name).await? {
1405            return Ok(None);
1406        }
1407
1408        // Filter columns to those that exist
1409        let actual_columns =
1410            if let Some(table_schema) = backend.get_table_schema(&table_name).await? {
1411                let table_field_names: HashSet<&str> = table_schema
1412                    .fields()
1413                    .iter()
1414                    .map(|f| f.name().as_str())
1415                    .collect();
1416                columns
1417                    .iter()
1418                    .copied()
1419                    .filter(|c| table_field_names.contains(c))
1420                    .map(|s| s.to_string())
1421                    .collect::<Vec<_>>()
1422            } else {
1423                return Ok(None);
1424            };
1425
1426        let filter = match (edge_hwm, additional_filter) {
1427            (Some(hwm), Some(f)) => Some(format!("_version <= {} AND ({})", hwm, f)),
1428            (Some(hwm), None) => Some(format!("_version <= {}", hwm)),
1429            (None, Some(f)) => Some(f.to_string()),
1430            (None, None) => None,
1431        };
1432
1433        let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1434        if let Some(f) = filter {
1435            request = request.with_filter(f);
1436        }
1437
1438        // Fail closed: a scan error (transient I/O, an unparsable filter, a
1439        // corrupt fragment) must propagate, never be silently mapped to
1440        // `Ok(None)`. Callers treat `Ok(None)` as "no rows" — e.g. the MERGE
1441        // fast path would create a duplicate node on a transient failure (review
1442        // bug #3a) — so an error here must surface as an error. A genuinely-
1443        // absent table is already handled above.
1444        let batches = backend.scan(request).await?;
1445        if batches.is_empty() {
1446            Ok(None)
1447        } else {
1448            Ok(Some(arrow::compute::concat_batches(
1449                &batches[0].schema(),
1450                &batches,
1451            )?))
1452        }
1453    }
1454
1455    /// Scan the unified main vertex table. Returns `None` if table doesn't exist.
1456    ///
1457    /// Applies version HWM filter internally for snapshot isolation, combined
1458    /// with any caller-provided filter (label conditions, etc.).
1459    pub async fn scan_main_vertex_table(
1460        &self,
1461        columns: &[&str],
1462        filter: Option<&str>,
1463    ) -> Result<Option<arrow_array::RecordBatch>> {
1464        let backend = self.backend();
1465        let table_name = table_names::main_vertex_table_name();
1466
1467        if !backend.table_exists(table_name).await? {
1468            return Ok(None);
1469        }
1470
1471        // Combine caller filter with version HWM for snapshot isolation
1472        let full_filter = match (self.version_high_water_mark(), filter) {
1473            (Some(hwm), Some(f)) => Some(format!("_version <= {} AND ({})", hwm, f)),
1474            (Some(hwm), None) => Some(format!("_version <= {}", hwm)),
1475            (None, Some(f)) => Some(f.to_string()),
1476            (None, None) => None,
1477        };
1478
1479        let request = ScanRequest::all(table_name)
1480            .with_columns(columns.iter().map(|s| s.to_string()).collect());
1481        let request = match full_filter.as_deref() {
1482            Some(f) => request.with_filter(f),
1483            None => request,
1484        };
1485
1486        // Fail closed: a scan error (transient I/O, an unparsable filter, a
1487        // corrupt fragment) must propagate, never be silently mapped to
1488        // `Ok(None)`. Callers treat `Ok(None)` as "no rows" — e.g. the MERGE
1489        // fast path would create a duplicate node on a transient failure (review
1490        // bug #3a) — so an error here must surface as an error. A genuinely-
1491        // absent table is already handled above.
1492        let batches = backend.scan(request).await?;
1493        if batches.is_empty() {
1494            Ok(None)
1495        } else {
1496            Ok(Some(arrow::compute::concat_batches(
1497                &batches[0].schema(),
1498                &batches,
1499            )?))
1500        }
1501    }
1502
1503    /// Scan the main edge table as a stream. Returns `None` if table doesn't exist.
1504    pub async fn scan_main_edge_table_stream(
1505        &self,
1506        filter: Option<&str>,
1507    ) -> Result<
1508        Option<
1509            std::pin::Pin<Box<dyn futures::Stream<Item = Result<arrow_array::RecordBatch>> + Send>>,
1510        >,
1511    > {
1512        let backend = self.backend();
1513        let table_name = table_names::main_edge_table_name();
1514
1515        if !backend.table_exists(table_name).await? {
1516            return Ok(None);
1517        }
1518
1519        let mut request = ScanRequest::all(table_name);
1520        if let Some(f) = filter {
1521            request = request.with_filter(f);
1522        }
1523
1524        let stream = backend.scan_stream(request).await?;
1525        Ok(Some(stream))
1526    }
1527
1528    /// Scan a per-label vertex table as a stream. Returns `None` if table doesn't exist.
1529    pub async fn scan_vertex_table_stream(
1530        &self,
1531        label: &str,
1532    ) -> Result<
1533        Option<
1534            std::pin::Pin<Box<dyn futures::Stream<Item = Result<arrow_array::RecordBatch>> + Send>>,
1535        >,
1536    > {
1537        let backend = self.backend();
1538        let table_name = table_names::vertex_table_name(label);
1539
1540        if !backend.table_exists(&table_name).await? {
1541            return Ok(None);
1542        }
1543
1544        let stream = backend.scan_stream(ScanRequest::all(&table_name)).await?;
1545        Ok(Some(stream))
1546    }
1547
1548    /// Find a vertex VID by external ID. Uses pinned snapshot HWM if present.
1549    pub async fn find_vertex_by_ext_id(&self, ext_id: &str) -> Result<Option<Vid>> {
1550        MainVertexDataset::find_by_ext_id(self.backend(), ext_id, self.version_high_water_mark())
1551            .await
1552    }
1553
1554    /// Map every live vertex that has an external id to its `ext_id`
1555    /// (`_vid` → `ext_id`).
1556    ///
1557    /// `ext_id` is folded into a vertex's content `_uid` but is stripped from
1558    /// query results, so the fork diff/promote engine can't recover it by
1559    /// re-hashing query rows — two vertices differing only by `ext_id` would
1560    /// collapse to one identity (review H4). This exposes the stored `ext_id`
1561    /// so the diff can fold it back into its recomputed UID. Reads through the
1562    /// (branched) backend, so a forked manager sees its own + inherited rows.
1563    /// Covers flushed (Lance) rows; `ext_id` is immutable so no version
1564    /// reconciliation is needed.
1565    pub async fn get_vertex_ext_ids(&self) -> Result<std::collections::HashMap<Vid, String>> {
1566        use arrow_array::StringArray;
1567        let backend = self.backend.as_ref();
1568        let vtable = table_names::main_vertex_table_name();
1569        let mut out = std::collections::HashMap::new();
1570        if !backend.table_exists(vtable).await? {
1571            return Ok(out);
1572        }
1573        let request = ScanRequest::all(vtable)
1574            .with_filter("_deleted = false")
1575            .with_columns(vec!["_vid".to_string(), "ext_id".to_string()]);
1576        let batches = backend
1577            .scan(request)
1578            .await
1579            .map_err(|e| anyhow!("get_vertex_ext_ids: {}", e))?;
1580        for batch in batches {
1581            let vids = batch
1582                .column_by_name("_vid")
1583                .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
1584                .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid _vid column"))?;
1585            let exts = batch
1586                .column_by_name("ext_id")
1587                .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1588                .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid ext_id column"))?;
1589            for i in 0..batch.num_rows() {
1590                if exts.is_null(i) {
1591                    continue;
1592                }
1593                let ext = exts.value(i);
1594                if !ext.is_empty() {
1595                    out.insert(Vid::from(vids.value(i)), ext.to_string());
1596                }
1597            }
1598        }
1599        Ok(out)
1600    }
1601
1602    /// Find labels for a vertex by VID. Uses pinned snapshot HWM if present.
1603    pub async fn find_vertex_labels_by_vid(&self, vid: Vid) -> Result<Option<Vec<String>>> {
1604        MainVertexDataset::find_labels_by_vid(self.backend(), vid, self.version_high_water_mark())
1605            .await
1606    }
1607
1608    /// Find edges from the main edge table by type names, optionally pushing
1609    /// a bounded endpoint vid set into the scan (review perf #5).
1610    pub async fn find_edges_by_type_names(
1611        &self,
1612        type_names: &[&str],
1613        endpoint_filter: Option<(crate::storage::main_edge::EndpointSide, &[Vid])>,
1614    ) -> Result<Vec<(Eid, Vid, Vid, String, uni_common::Properties)>> {
1615        MainEdgeDataset::find_edges_by_type_names(self.backend(), type_names, endpoint_filter).await
1616    }
1617
1618    /// Scan vertex candidates matching a filter. Returns VIDs where `_deleted = false`.
1619    pub async fn scan_vertex_candidates(
1620        &self,
1621        label: &str,
1622        filter: Option<&str>,
1623    ) -> Result<Vec<Vid>> {
1624        let backend = self.backend();
1625        let table_name = table_names::vertex_table_name(label);
1626
1627        if !backend.table_exists(&table_name).await? {
1628            return Ok(Vec::new());
1629        }
1630
1631        let full_filter = match filter {
1632            Some(f) => format!("_deleted = false AND ({})", f),
1633            None => "_deleted = false".to_string(),
1634        };
1635
1636        let request = ScanRequest::all(&table_name)
1637            .with_filter(full_filter)
1638            .with_columns(vec!["_vid".to_string()]);
1639
1640        let batches = backend.scan(request).await?;
1641
1642        let mut vids = Vec::new();
1643        for batch in batches {
1644            let vid_col = batch
1645                .column_by_name("_vid")
1646                .ok_or(anyhow!("Missing _vid"))?
1647                .as_any()
1648                .downcast_ref::<UInt64Array>()
1649                .ok_or(anyhow!("Invalid _vid"))?;
1650            for i in 0..batch.num_rows() {
1651                vids.push(Vid::from(vid_col.value(i)));
1652            }
1653        }
1654        Ok(vids)
1655    }
1656
1657    /// Construct a [`VertexDataset`] batch-builder for `label`.
1658    ///
1659    /// `VertexDataset` no longer opens on-disk data, so there is nothing to
1660    /// branch here — fork-scoped reads of vertex data go through the
1661    /// (branch-aware) `StorageBackend`.
1662    pub fn vertex_dataset(&self, label: &str) -> Result<VertexDataset> {
1663        let schema = self.schema_manager.schema();
1664        let label_meta = schema
1665            .labels
1666            .get(label)
1667            .ok_or_else(|| anyhow!("Label '{}' not found", label))?;
1668        Ok(VertexDataset::new(&self.base_uri, label, label_meta.id))
1669    }
1670
1671    #[cfg(feature = "lance-backend")]
1672    pub fn edge_dataset(
1673        &self,
1674        edge_type: &str,
1675        src_label: &str,
1676        dst_label: &str,
1677    ) -> Result<EdgeDataset> {
1678        let key = format!("edges_{edge_type}");
1679        match self.fork_branch_for(&key) {
1680            Some(branch) => Ok(EdgeDataset::new_branched(
1681                &self.base_uri,
1682                edge_type,
1683                src_label,
1684                dst_label,
1685                branch,
1686            )),
1687            None => Ok(EdgeDataset::new(
1688                &self.base_uri,
1689                edge_type,
1690                src_label,
1691                dst_label,
1692            )),
1693        }
1694    }
1695
1696    pub fn delta_dataset(&self, edge_type: &str, direction: &str) -> Result<DeltaDataset> {
1697        let key = format!("deltas_{edge_type}_{direction}");
1698        match self.fork_branch_for(&key) {
1699            Some(branch) => Ok(DeltaDataset::new_branched(
1700                &self.base_uri,
1701                edge_type,
1702                direction,
1703                branch,
1704            )),
1705            None => Ok(DeltaDataset::new(&self.base_uri, edge_type, direction)),
1706        }
1707    }
1708
1709    pub fn adjacency_dataset(
1710        &self,
1711        edge_type: &str,
1712        label: &str,
1713        direction: &str,
1714    ) -> Result<AdjacencyDataset> {
1715        // The fork registers adjacency branches under the canonical table
1716        // name (`adjacency_{edge_type}_{direction}`), so the lookup key must
1717        // match it — not the historical `adjacency_{direction}_{edge_type}_
1718        // {label}`, which never resolved a branch. Adjacency is per-`(edge_
1719        // type, direction)`, not per-label. The canonical form is pinned by
1720        // `table_names::tests::adjacency_table_name_is_canonical`. (L8)
1721        let key = crate::backend::table_names::adjacency_table_name(edge_type, direction);
1722        match self.fork_branch_for(&key) {
1723            Some(branch) => Ok(AdjacencyDataset::new_branched(
1724                &self.base_uri,
1725                edge_type,
1726                label,
1727                direction,
1728                branch,
1729            )),
1730            None => Ok(AdjacencyDataset::new(
1731                &self.base_uri,
1732                edge_type,
1733                label,
1734                direction,
1735            )),
1736        }
1737    }
1738
1739    /// Look up the branch name for a dataset under the active fork
1740    /// scope, if any. Returns `None` when not forked, or when the fork
1741    /// hasn't recorded a branch on this dataset yet (Phase 2 territory).
1742    fn fork_branch_for(&self, dataset_name: &str) -> Option<String> {
1743        self.fork_scope
1744            .as_ref()
1745            .and_then(|s| s.branch_for(dataset_name))
1746    }
1747
1748    /// Get the main vertex dataset for unified vertex storage.
1749    ///
1750    /// The main vertex dataset contains all vertices regardless of label,
1751    /// enabling fast ID-based lookups without knowing the label.
1752    pub fn main_vertex_dataset(&self) -> MainVertexDataset {
1753        MainVertexDataset::new(&self.base_uri)
1754    }
1755
1756    /// Get the main edge dataset for unified edge storage.
1757    ///
1758    /// The main edge dataset contains all edges regardless of type,
1759    /// enabling fast ID-based lookups without knowing the edge type.
1760    pub fn main_edge_dataset(&self) -> MainEdgeDataset {
1761        MainEdgeDataset::new(&self.base_uri)
1762    }
1763
1764    #[cfg(feature = "lance-backend")]
1765    pub fn uid_index(&self, label: &str) -> Result<UidIndex> {
1766        Ok(UidIndex::new(&self.base_uri, label))
1767    }
1768
1769    #[cfg(feature = "lance-backend")]
1770    pub async fn inverted_index(&self, label: &str, property: &str) -> Result<InvertedIndex> {
1771        let schema = self.schema_manager.schema();
1772        let config = schema
1773            .indexes
1774            .iter()
1775            .find_map(|idx| match idx {
1776                IndexDefinition::Inverted(cfg)
1777                    if cfg.label == label && cfg.property == property =>
1778                {
1779                    Some(cfg.clone())
1780                }
1781                _ => None,
1782            })
1783            .ok_or_else(|| anyhow!("Inverted index not found for {}.{}", label, property))?;
1784
1785        InvertedIndex::new(&self.base_uri, config).await
1786    }
1787
1788    #[expect(clippy::too_many_arguments)]
1789    pub async fn vector_search(
1790        &self,
1791        label: &str,
1792        property: &str,
1793        query: &[f32],
1794        k: usize,
1795        filter: Option<&str>,
1796        opts: VectorQueryOpts,
1797        ctx: Option<&QueryContext>,
1798    ) -> Result<Vec<(Vid, f32)>> {
1799        use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
1800
1801        // Look up vector index config to get the correct distance metric.
1802        let schema = self.schema_manager.schema();
1803
1804        // A query vector that doesn't match the declared column dimensions can never
1805        // score a row — every candidate would be skipped by the length guards below
1806        // and in the backend, silently returning 0 rows. Error instead (issue #137).
1807        // Undeclared (schemaless) properties skip the check: no dim is authoritative.
1808        if let Some(meta) = schema
1809            .properties
1810            .get(label)
1811            .and_then(|props| props.get(property))
1812            && let uni_common::core::schema::DataType::Vector { dimensions } = &meta.r#type
1813            && query.len() != *dimensions
1814        {
1815            return Err(anyhow!(
1816                "vector dimension mismatch: query vector has {} dimensions but '{}.{}' is \
1817                 declared VECTOR({})",
1818                query.len(),
1819                label,
1820                property,
1821                dimensions
1822            ));
1823        }
1824
1825        let metric = schema
1826            .vector_index_for_property(label, property)
1827            .map(|config| config.metric.clone())
1828            .unwrap_or(DistanceMetric::L2);
1829
1830        let backend = self.backend.as_ref();
1831        let name = table_names::vertex_table_name(label);
1832
1833        let mut results = Vec::new();
1834
1835        // Only search if the table exists
1836        if backend.table_exists(&name).await? {
1837            let backend_metric = match &metric {
1838                DistanceMetric::L2 => BackendMetric::L2,
1839                DistanceMetric::Cosine => BackendMetric::Cosine,
1840                DistanceMetric::Dot => BackendMetric::Dot,
1841                // L1 has no ANN backend metric (its column is never ANN-indexed);
1842                // candidates are fetched by L2 over the FULL table (`fetch_k`
1843                // below) then re-scored with exact L1, so this L2 ordering does
1844                // not affect the final result.
1845                DistanceMetric::L1 => BackendMetric::L2,
1846                _ => BackendMetric::L2,
1847            };
1848
1849            // Build combined filter: _deleted = false + optional user filter + HWM
1850            let mut filter_parts = vec![Self::build_active_filter(filter)];
1851            if ctx.is_some()
1852                && let Some(hwm) = self.version_high_water_mark()
1853            {
1854                filter_parts.push(format!("_version <= {}", hwm));
1855            }
1856            let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
1857
1858            // L1/Manhattan is served exact/brute-force: fetch every candidate
1859            // (limit = full row count) and rank by the exact L1 re-score below.
1860            // O(N) — the inherent cost of exact L1, which cannot use an ANN index.
1861            let fetch_k = if matches!(metric, DistanceMetric::L1) {
1862                backend.count_rows(&name, None).await.unwrap_or(k).max(k)
1863            } else {
1864                k
1865            };
1866
1867            let batches = backend
1868                .vector_search(
1869                    &name,
1870                    property,
1871                    query,
1872                    fetch_k,
1873                    backend_metric,
1874                    combined_filter,
1875                    opts,
1876                )
1877                .await?;
1878
1879            // Re-score ANN candidates with an EXACT distance rather than trusting
1880            // Lance's `_distance`. Lance's cosine ANN distance is on a different
1881            // scale (`2(1-cos)`) than the flat path (`1-cos`) that
1882            // `calculate_score` expects, so the raw value makes the final
1883            // similarity depend on whether an ANN index served the query
1884            // (issue #138). The candidate vectors come back in `batches`, so
1885            // re-scoring needs no extra fetch and puts these candidates on the
1886            // same `compute_distance` scale as the L0 candidates merged below.
1887            results = extract_vid_and_vector_pairs(&batches, "_vid", property)?
1888                .into_iter()
1889                .filter(|(_, emb)| emb.len() == query.len())
1890                .map(|(vid, emb)| (vid, metric.compute_distance(&emb, query)))
1891                .collect();
1892        }
1893
1894        // Merge L0 buffer vertices into results for visibility of unflushed data.
1895        if let Some(qctx) = ctx {
1896            merge_l0_into_vector_results(&mut results, qctx, label, property, query, k, &metric);
1897        }
1898
1899        // Exact top-k. The L1 over-fetch (and the no-L0-activity early return in
1900        // the merge) can leave more than `k` candidates; rank by distance
1901        // ascending and truncate. Idempotent for the ANN paths, which already
1902        // return `k` sorted.
1903        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1904        results.truncate(k);
1905
1906        Ok(results)
1907    }
1908
1909    /// First-stage candidate generation for a MUVERA index: single-vector ANN over the
1910    /// derived FDE column with the **Dot** metric (the FDE inner product approximates
1911    /// MaxSim — the physical index was built with Dot, see
1912    /// `IndexManager::create_vector_index`).
1913    ///
1914    /// Flushed/indexed data ONLY — unlike [`Self::vector_search`] this deliberately does
1915    /// NOT merge L0, because the live L0 has no FDE column (it is materialised at flush by
1916    /// `Writer::materialize_fde_columns`). L0 visibility is provided one layer up by
1917    /// `multivector_rerank`, which unions live L0 vids and re-scores everything by exact
1918    /// MaxSim. Scores returned here are placeholder distances (the caller re-ranks).
1919    #[expect(clippy::too_many_arguments)]
1920    pub async fn muvera_fde_candidates(
1921        &self,
1922        label: &str,
1923        fde_column: &str,
1924        fde_query: &[f32],
1925        k: usize,
1926        filter: Option<&str>,
1927        opts: VectorQueryOpts,
1928        ctx: Option<&QueryContext>,
1929    ) -> Result<Vec<(Vid, f32)>> {
1930        use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
1931
1932        let backend = self.backend.as_ref();
1933        let name = table_names::vertex_table_name(label);
1934        // Distinguish "table genuinely absent" (Ok(false) → nothing flushed yet, L0-only
1935        // candidates merged upstream) from a transient backend fault (Err): the latter
1936        // must surface, not silently degrade to incomplete results (issue #96).
1937        if !backend.table_exists(&name).await? {
1938            return Ok(Vec::new());
1939        }
1940
1941        let mut filter_parts = vec![Self::build_active_filter(filter)];
1942        if ctx.is_some()
1943            && let Some(hwm) = self.version_high_water_mark()
1944        {
1945            filter_parts.push(format!("_version <= {}", hwm));
1946        }
1947        let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
1948
1949        let batches = backend
1950            .vector_search(
1951                &name,
1952                fde_column,
1953                fde_query,
1954                k,
1955                BackendMetric::Dot,
1956                combined_filter,
1957                opts,
1958            )
1959            .await?;
1960        extract_vid_score_pairs(&batches, "_vid", "_distance")
1961    }
1962
1963    /// Late-interaction (ColBERT / MaxSim) first-stage search over a multi-vector
1964    /// (`List<Vector>`) column.
1965    ///
1966    /// Issues a multi-token query — Lance scores each row's token set by MaxSim —
1967    /// and defaults to **Cosine** (the ColBERT convention) when the property has
1968    /// no index, vs `L2` for dense vectors. `opts` (`nprobes` / `refine_factor`)
1969    /// tune the underlying ANN index.
1970    ///
1971    /// This is a **candidate generator over flushed/indexed data only** — it does
1972    /// not merge unflushed L0 rows (unlike [`Self::vector_search`], which can,
1973    /// because the single-vector in-process distance is on the identical scale as
1974    /// Lance's `_distance`). Lance's multi-vector `_distance` is an opaque internal
1975    /// aggregate whose scale cannot be matched against an in-process MaxSim, so L0
1976    /// visibility is provided one layer up: the uni-query `multivector_rerank`
1977    /// helper unions these flushed candidates with live L0 vids and re-scores
1978    /// *every* candidate by exact MaxSim. Callers wanting recent-write visibility
1979    /// must go through that path (the `uni.vector.query` procedure and the inline
1980    /// `vector_similarity` predicate both do); calling this directly sees flushed
1981    /// data only. Multi-vector search on forks/branches remains unsupported
1982    /// (`backend::branched` bails).
1983    #[expect(clippy::too_many_arguments)]
1984    pub async fn multivector_search(
1985        &self,
1986        label: &str,
1987        property: &str,
1988        query: &[Vec<f32>],
1989        k: usize,
1990        filter: Option<&str>,
1991        opts: VectorQueryOpts,
1992        ctx: Option<&QueryContext>,
1993    ) -> Result<Vec<(Vid, f32)>> {
1994        use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
1995
1996        let schema = self.schema_manager.schema();
1997        let metric = schema
1998            .vector_index_for_property(label, property)
1999            .map(|config| config.metric.clone())
2000            .unwrap_or(DistanceMetric::Cosine);
2001
2002        let backend = self.backend.as_ref();
2003        let name = table_names::vertex_table_name(label);
2004
2005        // On a branched table, Lance has no per-branch multi-vector nearest
2006        // (and lancedb cannot open a `Table` on a non-main branch), so the
2007        // backend's `multivector_search` bails. Instead, enumerate the branch's
2008        // candidate vids via a branch-aware scan (`BranchedBackend::scan`
2009        // applies the branch, surfacing fork-local + parent-inherited rows via
2010        // `base_paths`) and let the uni-query layer re-score by exact MaxSim —
2011        // it fetches candidate properties branch-aware and merges fork L0. The
2012        // returned score is a placeholder (the only caller re-ranks). This is a
2013        // brute-force scan, O(branch rows incl. inherited): the inherent cost of
2014        // having no multi-vector index on branches.
2015        let branched = self
2016            .fork_scope
2017            .as_ref()
2018            .is_some_and(|s| s.branch_for(&name).is_some());
2019        if branched {
2020            // Ok(false) = nothing flushed on the branch (fork L0 rows merged upstream);
2021            // Err = a backend fault that must surface rather than degrade silently (#96).
2022            if !backend.table_exists(&name).await? {
2023                return Ok(Vec::new());
2024            }
2025            let mut filter_parts = vec![Self::build_active_filter(filter)];
2026            if ctx.is_some()
2027                && let Some(hwm) = self.version_high_water_mark()
2028            {
2029                filter_parts.push(format!("_version <= {}", hwm));
2030            }
2031            let request = ScanRequest::all(&name)
2032                .with_filter(filter_parts.join(" AND "))
2033                .with_columns(vec!["_vid".to_string()]);
2034            let batches = backend.scan(request).await?;
2035            let mut results = Vec::new();
2036            for batch in batches {
2037                let vid_col = batch
2038                    .column_by_name("_vid")
2039                    .ok_or(anyhow!("Missing _vid"))?
2040                    .as_any()
2041                    .downcast_ref::<UInt64Array>()
2042                    .ok_or(anyhow!("Invalid _vid"))?;
2043                for i in 0..batch.num_rows() {
2044                    results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2045                }
2046            }
2047            return Ok(results);
2048        }
2049
2050        let mut results = Vec::new();
2051        // Ok(false) = no flushed table yet (L0-only); Err must propagate, not fail open (#96).
2052        if backend.table_exists(&name).await? {
2053            let backend_metric = match &metric {
2054                DistanceMetric::L2 => BackendMetric::L2,
2055                DistanceMetric::Cosine => BackendMetric::Cosine,
2056                DistanceMetric::Dot => BackendMetric::Dot,
2057                _ => BackendMetric::Cosine,
2058            };
2059
2060            let mut filter_parts = vec![Self::build_active_filter(filter)];
2061            if ctx.is_some()
2062                && let Some(hwm) = self.version_high_water_mark()
2063            {
2064                filter_parts.push(format!("_version <= {}", hwm));
2065            }
2066            let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
2067
2068            let batches = backend
2069                .multivector_search(
2070                    &name,
2071                    property,
2072                    query,
2073                    k,
2074                    backend_metric,
2075                    combined_filter,
2076                    opts,
2077                )
2078                .await?;
2079            results = extract_vid_score_pairs(&batches, "_vid", "_distance")?;
2080        }
2081
2082        Ok(results)
2083    }
2084
2085    /// Flushed-data candidate generation for scored sparse-vector retrieval.
2086    ///
2087    /// Loads the registered sparse index and returns its top-`k`
2088    /// `(vid, prelim_score)` by dot product over the flushed postings. Like
2089    /// [`Self::multivector_search`], this is **flushed-only** and the prelim
2090    /// score is advisory: the uni-query `sparse_rerank` helper unions live L0
2091    /// rows and re-scores *every* candidate exactly and MVCC-aware via
2092    /// `sparse_dot`, so callers wanting recent-write visibility must go through
2093    /// that path. Returns empty if no sparse index is registered for the
2094    /// property. On a fork/branch there is no per-branch sparse index, so this
2095    /// brute-force enumerates the branch's candidate vids (Approach A — see the
2096    /// branched arm) for the re-score path.
2097    pub async fn sparse_search(
2098        &self,
2099        label: &str,
2100        property: &str,
2101        query: &[(u32, f32)],
2102        k: usize,
2103    ) -> Result<Vec<(Vid, f32)>> {
2104        #[cfg(feature = "lance-backend")]
2105        {
2106            let name = table_names::vertex_table_name(label);
2107            let branched = self
2108                .fork_scope
2109                .as_ref()
2110                .is_some_and(|s| s.branch_for(&name).is_some());
2111            if branched {
2112                // Approach A (v1): the sparse index is a separate hand-rolled
2113                // Lance dataset, not a vertices-table index, so it cannot ride
2114                // Lance's `base_paths` branch fusion the way the dense/FTS
2115                // Lance-native indexes do — there is no per-branch sparse index to
2116                // query. Enumerate the branch's candidate vids via a branch-aware
2117                // scan (`base_paths` surfaces fork-local + parent-inherited rows;
2118                // the `_deleted = false` prefilter drops tombstoned inherited rows)
2119                // and let the uni-query `sparse_rerank` helper re-score every
2120                // candidate exactly via `sparse_dot` and union fork L0. The
2121                // returned score is a placeholder (the only caller re-ranks). This
2122                // is a brute-force scan, O(branch rows incl. inherited): the
2123                // proposal's brute-force-DAAT-first choice, mirroring
2124                // [`Self::multivector_search`]'s branched path.
2125                //
2126                // Approach B (deferred, benchmark-gated — issue #95 M5): build a
2127                // fork-local sparse postings dataset on a fork-scoped path
2128                // (`SparseVectorIndex::postings_path` made fork-aware), then query
2129                // parent ∪ fork-local postings minus the fork tombstone set,
2130                // resolving nested forks through ancestor postings datasets. Faster
2131                // on large fork corpora, but it re-implements by hand the fusion /
2132                // tombstone / nested-fork correctness this scan gets from Lance.
2133                let backend = self.backend.as_ref();
2134                // Ok(false) = nothing flushed on the branch (fork L0 rows merge upstream);
2135                // Err = a backend fault that must surface, not fail open silently (#95).
2136                if !backend.table_exists(&name).await? {
2137                    return Ok(Vec::new());
2138                }
2139                let request = ScanRequest::all(&name)
2140                    .with_filter(Self::build_active_filter(None))
2141                    .with_columns(vec!["_vid".to_string()]);
2142                let batches = backend.scan(request).await?;
2143                let mut results = Vec::new();
2144                for batch in batches {
2145                    let vid_col = batch
2146                        .column_by_name("_vid")
2147                        .ok_or_else(|| anyhow!("Missing _vid"))?
2148                        .as_any()
2149                        .downcast_ref::<UInt64Array>()
2150                        .ok_or_else(|| anyhow!("Invalid _vid"))?;
2151                    for i in 0..batch.num_rows() {
2152                        results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2153                    }
2154                }
2155                return Ok(results);
2156            }
2157            match self
2158                .index_manager()
2159                .sparse_vector_index(label, property)
2160                .await
2161            {
2162                Ok(idx) => idx.query_topk(query, k).await,
2163                Err(_) => Ok(Vec::new()),
2164            }
2165        }
2166        #[cfg(not(feature = "lance-backend"))]
2167        {
2168            let _ = (label, property, query, k);
2169            Ok(Vec::new())
2170        }
2171    }
2172
2173    /// Perform a full-text search with BM25 scoring.
2174    ///
2175    /// Returns vertices matching the search query along with their BM25 scores.
2176    /// Results are sorted by score descending (most relevant first).
2177    ///
2178    /// # Arguments
2179    /// * `label` - The label to search within
2180    /// * `property` - The property column to search (must have FTS index)
2181    /// * `query` - The search query text
2182    /// * `k` - Maximum number of results to return
2183    /// * `filter` - Optional Lance filter expression
2184    /// * `ctx` - Optional query context for visibility checks
2185    ///
2186    /// # Returns
2187    /// Vector of (Vid, score) tuples, where score is the BM25 relevance score.
2188    pub async fn fts_search(
2189        &self,
2190        label: &str,
2191        property: &str,
2192        query: &str,
2193        k: usize,
2194        filter: Option<&str>,
2195        ctx: Option<&QueryContext>,
2196    ) -> Result<Vec<(Vid, f32)>> {
2197        use crate::backend::types::FilterExpr;
2198
2199        let backend = self.backend.as_ref();
2200        let name = table_names::vertex_table_name(label);
2201
2202        let mut results = if backend.table_exists(&name).await? {
2203            // Build combined filter: _deleted = false + optional user filter + HWM
2204            let mut filter_parts = vec![Self::build_active_filter(filter)];
2205            if ctx.is_some()
2206                && let Some(hwm) = self.version_high_water_mark()
2207            {
2208                filter_parts.push(format!("_version <= {}", hwm));
2209            }
2210            let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
2211
2212            let batches = backend
2213                .full_text_search(&name, property, query, k, combined_filter)
2214                .await?;
2215
2216            let mut fts_results = extract_vid_score_pairs(&batches, "_vid", "_score")?;
2217            // Results should already be sorted by score from backend, but ensure descending order
2218            fts_results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2219            fts_results
2220        } else {
2221            Vec::new()
2222        };
2223
2224        // Merge L0 buffer vertices for visibility of unflushed data.
2225        if let Some(qctx) = ctx {
2226            merge_l0_into_fts_results(&mut results, qctx, label, property, query, k);
2227        }
2228
2229        Ok(results)
2230    }
2231
2232    #[cfg(feature = "lance-backend")]
2233    pub async fn get_vertex_by_uid(&self, uid: &UniId, label: &str) -> Result<Option<Vid>> {
2234        let index = self.uid_index(label)?;
2235        index.get_vid(uid).await
2236    }
2237
2238    #[cfg(feature = "lance-backend")]
2239    pub async fn insert_vertex_with_uid(&self, label: &str, vid: Vid, uid: UniId) -> Result<()> {
2240        let index = self.uid_index(label)?;
2241        index.write_mapping(&[(uid, vid)]).await
2242    }
2243
2244    pub async fn load_subgraph(
2245        &self,
2246        start_vids: &[Vid],
2247        edge_types: &[u32],
2248        max_hops: usize,
2249        direction: GraphDirection,
2250        l0: Option<&L0Buffer>,
2251    ) -> Result<WorkingGraph> {
2252        let mut graph = WorkingGraph::new();
2253        let schema = self.schema_manager.schema();
2254
2255        // Build maps for ID lookups
2256        let label_map: HashMap<u16, String> = schema
2257            .labels
2258            .values()
2259            .map(|meta| {
2260                (
2261                    meta.id,
2262                    schema.label_name_by_id(meta.id).unwrap().to_owned(),
2263                )
2264            })
2265            .collect();
2266
2267        let edge_type_map: HashMap<u32, String> = schema
2268            .edge_types
2269            .values()
2270            .map(|meta| {
2271                (
2272                    meta.id,
2273                    schema.edge_type_name_by_id(meta.id).unwrap().to_owned(),
2274                )
2275            })
2276            .collect();
2277
2278        let target_edge_types: HashSet<u32> = edge_types.iter().copied().collect();
2279
2280        // Initialize frontier
2281        let mut frontier: Vec<Vid> = start_vids.to_vec();
2282        let mut visited: HashSet<Vid> = HashSet::new();
2283
2284        // Add start vertices to graph
2285        for &vid in start_vids {
2286            graph.add_vertex(vid);
2287        }
2288
2289        for _hop in 0..max_hops {
2290            let mut next_frontier = HashSet::new();
2291
2292            for &vid in &frontier {
2293                if visited.contains(&vid) {
2294                    continue;
2295                }
2296                visited.insert(vid);
2297                graph.add_vertex(vid);
2298
2299                // For each edge type we want to traverse
2300                for &etype_id in &target_edge_types {
2301                    let etype_name = edge_type_map
2302                        .get(&etype_id)
2303                        .ok_or_else(|| anyhow!("Unknown edge type ID: {}", etype_id))?;
2304
2305                    // Determine directions
2306                    // Storage direction: "fwd" or "bwd".
2307                    // Query direction: Outgoing -> "fwd", Incoming -> "bwd".
2308                    let (dir_str, neighbor_is_dst) = match direction {
2309                        GraphDirection::Outgoing => ("fwd", true),
2310                        GraphDirection::Incoming => ("bwd", false),
2311                    };
2312
2313                    let mut edges: HashMap<Eid, EdgeState> = HashMap::new();
2314
2315                    // 1. L2: Adjacency (Base)
2316                    // In the new storage model, VIDs don't embed label info.
2317                    // We need to try all labels to find the adjacency data.
2318                    // Edge version from snapshot (reserved for future version filtering)
2319                    let _edge_ver = self
2320                        .pinned_snapshot
2321                        .as_ref()
2322                        .and_then(|s| s.edges.get(etype_name).map(|es| es.lance_version));
2323
2324                    // Try each label until we find adjacency data
2325                    let backend = self.backend();
2326                    for current_src_label in label_map.values() {
2327                        let adj_ds =
2328                            match self.adjacency_dataset(etype_name, current_src_label, dir_str) {
2329                                Ok(ds) => ds,
2330                                Err(_) => continue,
2331                            };
2332                        if let Some((neighbors, eids)) =
2333                            adj_ds.read_adjacency_backend(backend, vid).await?
2334                        {
2335                            for (n, eid) in neighbors.into_iter().zip(eids) {
2336                                edges.insert(
2337                                    eid,
2338                                    EdgeState {
2339                                        neighbor: n,
2340                                        version: 0,
2341                                        deleted: false,
2342                                    },
2343                                );
2344                            }
2345                            break; // Found adjacency data for this vid, no need to try other labels
2346                        }
2347                    }
2348
2349                    // 2. L1: Delta
2350                    let delta_ds = self.delta_dataset(etype_name, dir_str)?;
2351                    let delta_entries = delta_ds
2352                        .read_deltas(backend, vid, &schema, self.snapshot_version_hwm())
2353                        .await?;
2354                    Self::apply_delta_to_edges(&mut edges, delta_entries, neighbor_is_dst);
2355
2356                    // 3. L0: Buffer
2357                    if let Some(l0) = l0 {
2358                        Self::apply_l0_to_edges(&mut edges, l0, vid, etype_id, direction);
2359                    }
2360
2361                    // Add resulting edges to graph
2362                    Self::add_edges_to_graph(
2363                        &mut graph,
2364                        edges,
2365                        vid,
2366                        etype_id,
2367                        neighbor_is_dst,
2368                        &visited,
2369                        &mut next_frontier,
2370                    );
2371                }
2372            }
2373            frontier = next_frontier.into_iter().collect();
2374
2375            // Early termination: if frontier is empty, no more vertices to explore
2376            if frontier.is_empty() {
2377                break;
2378            }
2379        }
2380
2381        Ok(graph)
2382    }
2383
2384    /// Apply delta entries to edge state map, handling version conflicts.
2385    fn apply_delta_to_edges(
2386        edges: &mut HashMap<Eid, EdgeState>,
2387        delta_entries: Vec<crate::storage::delta::L1Entry>,
2388        neighbor_is_dst: bool,
2389    ) {
2390        for entry in delta_entries {
2391            let neighbor = if neighbor_is_dst {
2392                entry.dst_vid
2393            } else {
2394                entry.src_vid
2395            };
2396            let current_ver = edges.get(&entry.eid).map(|s| s.version).unwrap_or(0);
2397
2398            if entry.version > current_ver {
2399                edges.insert(
2400                    entry.eid,
2401                    EdgeState {
2402                        neighbor,
2403                        version: entry.version,
2404                        deleted: matches!(entry.op, Op::Delete),
2405                    },
2406                );
2407            }
2408        }
2409    }
2410
2411    /// Apply L0 buffer edges and tombstones to edge state map.
2412    fn apply_l0_to_edges(
2413        edges: &mut HashMap<Eid, EdgeState>,
2414        l0: &L0Buffer,
2415        vid: Vid,
2416        etype_id: u32,
2417        direction: GraphDirection,
2418    ) {
2419        let l0_neighbors = l0.get_neighbors(vid, etype_id, direction);
2420        for (neighbor, eid, ver) in l0_neighbors {
2421            let current_ver = edges.get(&eid).map(|s| s.version).unwrap_or(0);
2422            if ver > current_ver {
2423                edges.insert(
2424                    eid,
2425                    EdgeState {
2426                        neighbor,
2427                        version: ver,
2428                        deleted: false,
2429                    },
2430                );
2431            }
2432        }
2433
2434        // Check tombstones in L0
2435        for (eid, state) in edges.iter_mut() {
2436            if l0.is_tombstoned(*eid) {
2437                state.deleted = true;
2438            }
2439        }
2440    }
2441
2442    /// Add non-deleted edges to graph and collect next frontier.
2443    fn add_edges_to_graph(
2444        graph: &mut WorkingGraph,
2445        edges: HashMap<Eid, EdgeState>,
2446        vid: Vid,
2447        etype_id: u32,
2448        neighbor_is_dst: bool,
2449        visited: &HashSet<Vid>,
2450        next_frontier: &mut HashSet<Vid>,
2451    ) {
2452        for (eid, state) in edges {
2453            if state.deleted {
2454                continue;
2455            }
2456            graph.add_vertex(state.neighbor);
2457
2458            if !visited.contains(&state.neighbor) {
2459                next_frontier.insert(state.neighbor);
2460            }
2461
2462            if neighbor_is_dst {
2463                graph.add_edge(vid, state.neighbor, eid, etype_id);
2464            } else {
2465                graph.add_edge(state.neighbor, vid, eid, etype_id);
2466            }
2467        }
2468    }
2469}
2470
2471/// Extracts `(Vid, f32)` pairs from record batches using the given VID and score column names.
2472fn extract_vid_score_pairs(
2473    batches: &[arrow_array::RecordBatch],
2474    vid_column: &str,
2475    score_column: &str,
2476) -> Result<Vec<(Vid, f32)>> {
2477    let mut results = Vec::new();
2478    for batch in batches {
2479        let vid_col = batch
2480            .column_by_name(vid_column)
2481            .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2482            .as_any()
2483            .downcast_ref::<UInt64Array>()
2484            .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2485
2486        let score_col = batch
2487            .column_by_name(score_column)
2488            .ok_or_else(|| anyhow!("Missing {} column", score_column))?
2489            .as_any()
2490            .downcast_ref::<Float32Array>()
2491            .ok_or_else(|| anyhow!("Invalid {} column type", score_column))?;
2492
2493        for i in 0..batch.num_rows() {
2494            results.push((Vid::from(vid_col.value(i)), score_col.value(i)));
2495        }
2496    }
2497    Ok(results)
2498}
2499
2500/// Extracts `(vid, embedding)` pairs from vector-search result batches, decoding
2501/// the dense `FixedSizeList<Float32>` vector column.
2502///
2503/// Used to re-score ANN candidates with an exact `compute_distance` (issue #138):
2504/// Lance's cosine `_distance` scale differs between the ANN-index path
2505/// (`2(1-cos)`) and the flat path (`1-cos`), so the raw value cannot be trusted
2506/// for the final similarity. Rows whose vector is null or the wrong length are
2507/// skipped by the caller.
2508fn extract_vid_and_vector_pairs(
2509    batches: &[arrow_array::RecordBatch],
2510    vid_column: &str,
2511    vector_column: &str,
2512) -> Result<Vec<(Vid, Vec<f32>)>> {
2513    use arrow_array::{Array, FixedSizeListArray};
2514    let mut out = Vec::new();
2515    for batch in batches {
2516        let vid_col = batch
2517            .column_by_name(vid_column)
2518            .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2519            .as_any()
2520            .downcast_ref::<UInt64Array>()
2521            .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2522
2523        let vec_col = batch
2524            .column_by_name(vector_column)
2525            .ok_or_else(|| anyhow!("Missing vector column {}", vector_column))?
2526            .as_any()
2527            .downcast_ref::<FixedSizeListArray>()
2528            .ok_or_else(|| anyhow!("Vector column {} is not FixedSizeList", vector_column))?;
2529
2530        for i in 0..batch.num_rows() {
2531            if vec_col.is_null(i) {
2532                continue;
2533            }
2534            let element = vec_col.value(i);
2535            let floats = element
2536                .as_any()
2537                .downcast_ref::<Float32Array>()
2538                .ok_or_else(|| {
2539                    anyhow!("Vector column {} inner type is not Float32", vector_column)
2540                })?;
2541            let emb: Vec<f32> = (0..floats.len()).map(|j| floats.value(j)).collect();
2542            out.push((Vid::from(vid_col.value(i)), emb));
2543        }
2544    }
2545    Ok(out)
2546}
2547
2548/// Extracts a dense `Vec<f32>` embedding from an L0 property value.
2549///
2550/// Accepts both representations of a dense vector: the typed
2551/// [`Value::Vector`] (what the Cypher write path stores) and a
2552/// [`Value::List`] of numbers (the JSON-ingest representation), via the
2553/// canonical `TryFrom<&Value>` converter. Returns `None` if the property is
2554/// missing or not coercible to a numeric vector.
2555///
2556/// # Why both variants
2557///
2558/// Real Cypher writes land dense embeddings in L0 as `Value::Vector`; scoring
2559/// L0 candidates off only `Value::List` silently dropped every such candidate,
2560/// so committed-but-unflushed inserts/updates were invisible to dense search.
2561fn extract_embedding_from_props(
2562    props: &uni_common::Properties,
2563    property: &str,
2564) -> Option<Vec<f32>> {
2565    Vec::<f32>::try_from(props.get(property)?).ok()
2566}
2567
2568/// Merges L0 buffer vertices into LanceDB vector search results.
2569///
2570/// Visits L0 buffers in precedence order (pending flush → main → transaction),
2571/// collects tombstoned VIDs and candidate embeddings, then merges them with the
2572/// existing LanceDB results so that:
2573/// - Tombstoned VIDs are removed (unless re-created in a later L0).
2574/// - VIDs present in both L0 and LanceDB use the L0 distance.
2575/// - New L0-only VIDs are appended.
2576/// - Results are re-sorted by distance ascending and truncated to `k`.
2577fn merge_l0_into_vector_results(
2578    results: &mut Vec<(Vid, f32)>,
2579    ctx: &QueryContext,
2580    label: &str,
2581    property: &str,
2582    query: &[f32],
2583    k: usize,
2584    metric: &DistanceMetric,
2585) {
2586    // Collect all L0 buffers in precedence order (earliest first, last writer wins).
2587    let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2588        ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2589    buffers.push(Arc::clone(&ctx.l0));
2590    if let Some(ref txn) = ctx.transaction_l0 {
2591        buffers.push(Arc::clone(txn));
2592    }
2593
2594    // Maps VID → distance for L0 candidates (last writer wins).
2595    let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2596    // Tombstoned VIDs across all L0 buffers.
2597    let mut tombstoned: HashSet<Vid> = HashSet::new();
2598
2599    for buf_arc in &buffers {
2600        let buf = buf_arc.read();
2601
2602        // Accumulate tombstones.
2603        for &vid in &buf.vertex_tombstones {
2604            tombstoned.insert(vid);
2605        }
2606
2607        // Scan vertices with the target label.
2608        for (&vid, labels) in &buf.vertex_labels {
2609            if !labels.iter().any(|l| l == label) {
2610                continue;
2611            }
2612            if let Some(props) = buf.vertex_properties.get(&vid)
2613                && let Some(emb) = extract_embedding_from_props(props, property)
2614            {
2615                if emb.len() != query.len() {
2616                    // Inner defense only: `vector_search` rejects wrong-dim queries
2617                    // against declared columns up front, and the write paths reject
2618                    // wrong-dim values (issue #137), so for post-#137 data this skip
2619                    // is unreachable. It still guards mixed-dim rows written by older
2620                    // versions and undeclared (schemaless) properties, where
2621                    // `compute_distance` would panic on a length mismatch.
2622                    continue; // dimension mismatch
2623                }
2624                let dist = metric.compute_distance(&emb, query);
2625                // Last writer wins: later buffer overwrites earlier.
2626                l0_candidates.insert(vid, dist);
2627                // If re-created in a later L0, remove from tombstones.
2628                tombstoned.remove(&vid);
2629            }
2630        }
2631    }
2632
2633    // If no L0 activity affects this search, skip merge.
2634    if l0_candidates.is_empty() && tombstoned.is_empty() {
2635        return;
2636    }
2637
2638    // Remove tombstoned VIDs from LanceDB results.
2639    results.retain(|(vid, _)| !tombstoned.contains(vid));
2640
2641    // Overwrite or append L0 candidates.
2642    for (vid, dist) in &l0_candidates {
2643        // Skip a vid the precedence chain ultimately tombstoned: it can be a
2644        // live candidate from an earlier buffer while a later buffer deleted it
2645        // (that later buffer never revisits it in the label loop, so it stays
2646        // in `l0_candidates`). Appending it here would resurrect a deleted
2647        // vertex into vector-search results.
2648        if tombstoned.contains(vid) {
2649            continue;
2650        }
2651        if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2652            existing.1 = *dist;
2653        } else {
2654            results.push((*vid, *dist));
2655        }
2656    }
2657
2658    // Re-sort by distance ascending.
2659    results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
2660    results.truncate(k);
2661}
2662
2663/// Collects the live L0 vids carrying `label`, plus the set of vids tombstoned,
2664/// across the 3-tier L0 chain (pending flush → main → transaction).
2665///
2666/// Walks the buffers in precedence order (last writer wins): a vid created in a
2667/// later buffer clears an earlier tombstone, and a vid tombstoned in a later
2668/// buffer is removed from the live set. The returned `live` set therefore
2669/// excludes anything currently tombstoned.
2670///
2671/// This mirrors the L0 traversal in `merge_l0_into_vector_results` but returns
2672/// membership only — it does not score — so callers that re-score candidates
2673/// themselves (e.g. multi-vector MaxSim re-ranking, where Lance's `_distance`
2674/// scale is opaque) can build a candidate set without duplicating the
2675/// tombstone/precedence semantics.
2676pub fn collect_l0_label_candidates(ctx: &QueryContext, label: &str) -> (Vec<Vid>, HashSet<Vid>) {
2677    // Buffers in precedence order: pending flush → main → transaction.
2678    let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2679        ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2680    buffers.push(Arc::clone(&ctx.l0));
2681    if let Some(ref txn) = ctx.transaction_l0 {
2682        buffers.push(Arc::clone(txn));
2683    }
2684
2685    let mut live: HashSet<Vid> = HashSet::new();
2686    let mut tombstoned: HashSet<Vid> = HashSet::new();
2687
2688    for buf_arc in &buffers {
2689        let buf = buf_arc.read();
2690
2691        // A delete in this buffer wins over earlier creations.
2692        for &vid in &buf.vertex_tombstones {
2693            tombstoned.insert(vid);
2694            live.remove(&vid);
2695        }
2696
2697        // A (re-)creation with the target label in this buffer wins over an
2698        // earlier tombstone.
2699        for (&vid, labels) in &buf.vertex_labels {
2700            if !labels.iter().any(|l| l == label) {
2701                continue;
2702            }
2703            if buf.vertex_properties.contains_key(&vid) {
2704                live.insert(vid);
2705                tombstoned.remove(&vid);
2706            }
2707        }
2708    }
2709
2710    (live.into_iter().collect(), tombstoned)
2711}
2712
2713/// Computes a simple token-overlap relevance score between a query and text.
2714///
2715/// Returns the fraction of query tokens found in the text (case-insensitive),
2716/// producing a score in [0.0, 1.0]. Sufficient for the small L0 buffer.
2717fn compute_text_relevance(query: &str, text: &str) -> f32 {
2718    let query_tokens: HashSet<String> =
2719        query.split_whitespace().map(|t| t.to_lowercase()).collect();
2720    if query_tokens.is_empty() {
2721        return 0.0;
2722    }
2723    let text_tokens: HashSet<String> = text.split_whitespace().map(|t| t.to_lowercase()).collect();
2724    let hits = query_tokens
2725        .iter()
2726        .filter(|t| text_tokens.contains(t.as_str()))
2727        .count();
2728    hits as f32 / query_tokens.len() as f32
2729}
2730
2731/// Extracts a string slice from a property value.
2732fn extract_text_from_props<'a>(
2733    props: &'a uni_common::Properties,
2734    property: &str,
2735) -> Option<&'a str> {
2736    props.get(property)?.as_str()
2737}
2738
2739/// Merges L0 buffer vertices into LanceDB full-text search results.
2740///
2741/// Follows the same pattern as [`merge_l0_into_vector_results`]: visits L0
2742/// buffers in precedence order, collects tombstoned VIDs and text-match
2743/// candidates, then merges them so that:
2744/// - Tombstoned VIDs are removed (unless re-created in a later L0).
2745/// - VIDs present in both L0 and LanceDB use the L0 score.
2746/// - New L0-only VIDs are appended.
2747/// - Results are re-sorted by score **descending** and truncated to `k`.
2748fn merge_l0_into_fts_results(
2749    results: &mut Vec<(Vid, f32)>,
2750    ctx: &QueryContext,
2751    label: &str,
2752    property: &str,
2753    query: &str,
2754    k: usize,
2755) {
2756    // Collect all L0 buffers in precedence order (earliest first, last writer wins).
2757    let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2758        ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2759    buffers.push(Arc::clone(&ctx.l0));
2760    if let Some(ref txn) = ctx.transaction_l0 {
2761        buffers.push(Arc::clone(txn));
2762    }
2763
2764    // Maps VID → relevance score for L0 candidates (last writer wins).
2765    let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2766    // Tombstoned VIDs across all L0 buffers.
2767    let mut tombstoned: HashSet<Vid> = HashSet::new();
2768
2769    for buf_arc in &buffers {
2770        let buf = buf_arc.read();
2771
2772        // Accumulate tombstones.
2773        for &vid in &buf.vertex_tombstones {
2774            tombstoned.insert(vid);
2775        }
2776
2777        // Scan vertices with the target label.
2778        for (&vid, labels) in &buf.vertex_labels {
2779            if !labels.iter().any(|l| l == label) {
2780                continue;
2781            }
2782            if let Some(props) = buf.vertex_properties.get(&vid)
2783                && let Some(text) = extract_text_from_props(props, property)
2784            {
2785                let score = compute_text_relevance(query, text);
2786                if score > 0.0 {
2787                    // Last writer wins: later buffer overwrites earlier.
2788                    l0_candidates.insert(vid, score);
2789                }
2790                // If re-created in a later L0, remove from tombstones.
2791                tombstoned.remove(&vid);
2792            }
2793        }
2794    }
2795
2796    // If no L0 activity affects this search, skip merge.
2797    if l0_candidates.is_empty() && tombstoned.is_empty() {
2798        return;
2799    }
2800
2801    // Remove tombstoned VIDs from LanceDB results.
2802    results.retain(|(vid, _)| !tombstoned.contains(vid));
2803
2804    // Overwrite or append L0 candidates.
2805    for (vid, score) in &l0_candidates {
2806        // Skip a vid the precedence chain ultimately tombstoned (see the vector
2807        // helper above): a later buffer's delete must not be resurrected into
2808        // full-text-search results by an earlier buffer's live candidate.
2809        if tombstoned.contains(vid) {
2810            continue;
2811        }
2812        if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2813            existing.1 = *score;
2814        } else {
2815            results.push((*vid, *score));
2816        }
2817    }
2818
2819    // Re-sort by score descending (higher relevance first).
2820    results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2821    results.truncate(k);
2822}