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