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