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