From d502603a525c5e5093d042c97cd74f7fed939614 Mon Sep 17 00:00:00 2001 From: aaltshuler Date: Sat, 4 Jul 2026 18:04:38 +0300 Subject: [PATCH] perf(engine): attach the shared per-graph Session to every write/maintenance open MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Thread the graph's one Lance Session (previously read-path-only via ReadCaches) into TableStore, so open_dataset_head, the branch ops, and open_at_entry all attach it through the unified opener — the read path's handle cache and the write side now warm the same Lance metadata/index caches. diff_snapshots takes the graph's &TableStore instead of constructing its own session-less store from a root uri. Measured effect (write_cost.rs): the local data-table SCAN term — the merge-insert/RI scan that re-read O(depth) immutable fragment metadata per write — collapses from growing-with-depth to flat (depth 10: 1 read, depth 100: 1 read). The opener/scan-split gate is re-pinned to the new behavior and renamed (data_table_reads_split_into_flat_opener_and_scan_flat_with_session): a red there now means a write-side open dropped the session. The Snapshot-without-caches fallback stays session-less by design (a detached snapshot has no graph to share a session with); the S3 acceptance of the opener term remains owed to write_cost_s3 per the RFC-013 handoff. testing.md's backend-split note updated in the same change (the "local scan grows with depth" claim is now stale). --- crates/omnigraph/src/changes/mod.rs | 11 +++-- crates/omnigraph/src/db/manifest.rs | 10 +++-- crates/omnigraph/src/db/manifest/recovery.rs | 10 ++++- crates/omnigraph/src/db/manifest/tests.rs | 2 +- crates/omnigraph/src/db/omnigraph.rs | 45 +++++++++++--------- crates/omnigraph/src/instrumentation.rs | 12 +++--- crates/omnigraph/src/table_store.rs | 18 +++++--- crates/omnigraph/tests/failpoints.rs | 2 +- crates/omnigraph/tests/recovery.rs | 4 +- crates/omnigraph/tests/staged_writes.rs | 38 ++++++++--------- crates/omnigraph/tests/write_cost.rs | 28 +++++++----- crates/omnigraph/tests/writes.rs | 2 +- docs/dev/testing.md | 2 +- 13 files changed, 107 insertions(+), 77 deletions(-) diff --git a/crates/omnigraph/src/changes/mod.rs b/crates/omnigraph/src/changes/mod.rs index 2e9bc029..8ced466a 100644 --- a/crates/omnigraph/src/changes/mod.rs +++ b/crates/omnigraph/src/changes/mod.rs @@ -111,13 +111,12 @@ impl ChangeFilter { /// 2. Lineage check — same branch → version-column diff; different → ID-based diff /// 3. Row-level diff pub async fn diff_snapshots( - root_uri: &str, + table_store: &TableStore, from: &Snapshot, to: &Snapshot, filter: &ChangeFilter, branch: Option, ) -> Result { - let table_store = TableStore::new(root_uri); let mut all_keys: HashSet = HashSet::new(); for entry in from.entries() { all_keys.insert(entry.table_key.clone()); @@ -146,14 +145,14 @@ pub async fn diff_snapshots( let table_changes = if from_entry.is_none() { // Table added — all rows are inserts - diff_table_added(&table_store, to, table_key, is_edge, filter).await? + diff_table_added(table_store, to, table_key, is_edge, filter).await? } else if to_entry.is_none() { // Table removed — all rows are deletes - diff_table_removed(&table_store, from, table_key, is_edge, filter).await? + diff_table_removed(table_store, from, table_key, is_edge, filter).await? } else if same_lineage(from_entry, to_entry) { // Fast path: version-column diff diff_table_same_lineage( - &table_store, + table_store, from_entry.unwrap(), to_entry.unwrap(), is_edge, @@ -162,7 +161,7 @@ pub async fn diff_snapshots( .await? } else { // Cross-branch path: streaming ID-based diff - diff_table_cross_branch(&table_store, from, to, table_key, is_edge, filter).await? + diff_table_cross_branch(table_store, from, to, table_key, is_edge, filter).await? }; for mut c in table_changes { diff --git a/crates/omnigraph/src/db/manifest.rs b/crates/omnigraph/src/db/manifest.rs index f9bfdfb8..9b9853d2 100644 --- a/crates/omnigraph/src/db/manifest.rs +++ b/crates/omnigraph/src/db/manifest.rs @@ -159,7 +159,7 @@ impl Snapshot { ) .await } - None => entry.open(&self.root_uri).await, + None => entry.open(&self.root_uri, None).await, } } @@ -247,7 +247,11 @@ impl SubTableEntry { /// already holds the path, version, and branch. Branches are Lance native /// branches, so `with_branch` resolves `{base}/tree/{branch}` from the base /// URI; main uses `with_version`. - pub(crate) async fn open(&self, root_uri: &str) -> Result { + pub(crate) async fn open( + &self, + root_uri: &str, + session: Option<&Arc>, + ) -> Result { // The branch-qualified location is the dataset that physically holds this // version: main at `{table_path}`, a branch at // `{table_path}/tree/{branch}` (Lance native-branch storage). `with_version` @@ -265,7 +269,7 @@ impl SubTableEntry { crate::instrumentation::open_dataset( &location, crate::instrumentation::VersionResolution::At(self.table_version), - None, + session, crate::instrumentation::table_wrapper(), ) .await diff --git a/crates/omnigraph/src/db/manifest/recovery.rs b/crates/omnigraph/src/db/manifest/recovery.rs index 2c224a63..fbc02be1 100644 --- a/crates/omnigraph/src/db/manifest/recovery.rs +++ b/crates/omnigraph/src/db/manifest/recovery.rs @@ -2282,7 +2282,10 @@ mod tests { async fn restore_table_to_version_appends_one_commit() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new( + dir.path().to_str().unwrap(), + Arc::new(lance::session::Session::default()), + ); let mut ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) .await @@ -2320,7 +2323,10 @@ mod tests { // Repeated restore calls each produce a new HEAD+1 commit. let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new( + dir.path().to_str().unwrap(), + Arc::new(lance::session::Session::default()), + ); let mut ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) .await diff --git a/crates/omnigraph/src/db/manifest/tests.rs b/crates/omnigraph/src/db/manifest/tests.rs index 39206c22..0c87c5ca 100644 --- a/crates/omnigraph/src/db/manifest/tests.rs +++ b/crates/omnigraph/src/db/manifest/tests.rs @@ -154,7 +154,7 @@ async fn test_commit_changes_can_register_new_table_and_tombstone_old_one() { let ds = crate::table_store::TableStore::create_empty_dataset(&dataset_uri, &schema) .await .unwrap(); - let state = crate::table_store::TableStore::new(uri) + let state = crate::table_store::TableStore::new(uri, Arc::new(lance::session::Session::default())) .table_state(&dataset_uri, &ds) .await .unwrap(); diff --git a/crates/omnigraph/src/db/omnigraph.rs b/crates/omnigraph/src/db/omnigraph.rs index 48440356..c3bfc2da 100644 --- a/crates/omnigraph/src/db/omnigraph.rs +++ b/crates/omnigraph/src/db/omnigraph.rs @@ -86,9 +86,9 @@ pub struct SchemaApplyPreview { /// (when `base` is captured). NOT a general "no re-resolution" handle — the /// commit-time OCC re-read, the live-HEAD drift probe, and the fork-authority reads /// stay fresh (correctness machinery). Step 5 (PublishPlan unification) makes this -/// the non-optional publish carrier and adds session-aware base opens there, gated -/// by an S3 cost test — the warm-session benefit on the single remaining open is an -/// object-store phenomenon, so it earns its own gate rather than riding this PR. +/// the non-optional publish carrier. (Write/maintenance opens attach the shared +/// per-graph `Session` via the `TableStore`-held handle — the dataset-opener +/// unification; the S3 cost gate for that term is still owed.) /// /// Threaded as `Option<&WriteTxn>` through the mutate/load write chain /// (`open_for_mutation_on_branch`, `commit_all`, `commit_updates_on_branch_with_expected`) @@ -356,18 +356,21 @@ impl Omnigraph { } }; + let session = Arc::new(lance::session::Session::default()); Ok(Self { root_uri: root.clone(), storage, coordinator: Arc::new(tokio::sync::RwLock::new(coordinator)), - table_store: TableStore::new(&root), - runtime_cache: RuntimeCache::default(), // One shared Session per graph (LanceDB's one-session-per-connection - // model) plus the held-handle cache, created once and reused across - // reads. Session::default() caps are lazy (6 GiB index / 1 GiB - // metadata); multi-graph cap/sharing is a deferred follow-up. + // model), created once and shared by the read path (handle cache) + // AND the TableStore's write/maintenance opens, so both sides warm + // the same Lance metadata/index caches. Session::default() caps are + // lazy (6 GiB index / 1 GiB metadata); multi-graph cap/sharing is a + // deferred follow-up. + table_store: TableStore::new(&root, session.clone()), + runtime_cache: RuntimeCache::default(), read_caches: Arc::new(crate::runtime_cache::ReadCaches { - session: Arc::new(lance::session::Session::default()), + session, handles: Arc::new(crate::runtime_cache::TableHandleCache::default()), }), catalog: Arc::new(ArcSwap::from_pointee(catalog)), @@ -460,18 +463,21 @@ impl Omnigraph { let mut catalog = build_catalog_from_ir(&accepted_ir)?; fixup_blob_schemas(&mut catalog); + let session = Arc::new(lance::session::Session::default()); Ok(Self { root_uri: root.clone(), storage, coordinator: Arc::new(tokio::sync::RwLock::new(coordinator)), - table_store: TableStore::new(&root), - runtime_cache: RuntimeCache::default(), // One shared Session per graph (LanceDB's one-session-per-connection - // model) plus the held-handle cache, created once and reused across - // reads. Session::default() caps are lazy (6 GiB index / 1 GiB - // metadata); multi-graph cap/sharing is a deferred follow-up. + // model), created once and shared by the read path (handle cache) + // AND the TableStore's write/maintenance opens, so both sides warm + // the same Lance metadata/index caches. Session::default() caps are + // lazy (6 GiB index / 1 GiB metadata); multi-graph cap/sharing is a + // deferred follow-up. + table_store: TableStore::new(&root, session.clone()), + runtime_cache: RuntimeCache::default(), read_caches: Arc::new(crate::runtime_cache::ReadCaches { - session: Arc::new(lance::session::Session::default()), + session, handles: Arc::new(crate::runtime_cache::TableHandleCache::default()), }), catalog: Arc::new(ArcSwap::from_pointee(catalog)), @@ -776,8 +782,9 @@ impl Omnigraph { /// which still resolve through `snapshot_for_branch` and re-validate. Those reads must /// observe LIVE committed state, so unifying them (validate-once + pinned + re-checked /// read-set) is step 4's §7.1 work — threading `txn.base` there would re-introduce the - /// stale-read class the #298 cardinality fix removed. A session-aware base open is - /// likewise deferred to step 5 (handoff §1d). + /// stale-read class the #298 cardinality fix removed. Write-side opens now attach the shared + /// per-graph `Session` (the dataset-opener unification); the S3 cost gate + /// for that term is still owed (handoff §1d). pub(crate) async fn open_write_txn(&self, branch: Option<&str>) -> Result { let resolved = self.resolved_branch_target(branch).await?; Ok(WriteTxn { @@ -1195,7 +1202,7 @@ impl Omnigraph { let from_resolved = self.resolved_target(from).await?; let to_resolved = self.resolved_target(to).await?; crate::changes::diff_snapshots( - self.uri(), + &self.table_store, &from_resolved.snapshot, &to_resolved.snapshot, filter, @@ -1229,7 +1236,7 @@ impl Omnigraph { .await?; drop(coord); crate::changes::diff_snapshots( - self.uri(), + &self.table_store, &from_snap.snapshot, &to_snap.snapshot, filter, diff --git a/crates/omnigraph/src/instrumentation.rs b/crates/omnigraph/src/instrumentation.rs index 0f693d29..f4ef9ef9 100644 --- a/crates/omnigraph/src/instrumentation.rs +++ b/crates/omnigraph/src/instrumentation.rs @@ -42,8 +42,8 @@ pub struct QueryIoProbes { /// handle cache (Fix 3) serves them. pub table_wrapper: Option>, pub probe_count: Arc, - /// Counts DATA-table open CALLS through the two instrumented chokepoints - /// (`open_dataset_tracked` / `open_table_dataset`), classified by URI so the + /// Counts DATA-table open CALLS through the one instrumented chokepoint + /// (`open_dataset`), classified by URI so the /// internal/system tables (`__manifest`) are EXCLUDED — the publisher CAS /// opens those every write, and counting them would make the /// `data_open_count <= |touched_tables|` write gate @@ -52,9 +52,9 @@ pub struct QueryIoProbes { /// an exact open-invocation count. `forbidden_apis` keeps engine code OUTSIDE the /// storage layer (`exec/`, `db/omnigraph/`, `loader/`, `changes/`) from opening /// datasets except through these chokepoints, so the count is complete for the - /// keyed-write data path the gate measures. (`table_store.rs` is allow-listed and - /// does hold direct `Dataset::open`s — but only for branch-management ops - /// (`delete_branch`/`list_branches`/`force_delete_branch`), never that hot path.) + /// keyed-write data path the gate measures. (Since the dataset-opener + /// unification, `table_store.rs`'s branch-management ops also route through + /// the one chokepoint, so the count covers them too.) pub data_open_count: Arc, /// Internal/system-table (`__manifest`) open CALLS — the complement of /// `data_open_count`, kept for symmetry and debugging. @@ -142,7 +142,7 @@ fn open_is_internal(uri: &str) -> bool { /// table class (the URI's last segment) so the write gate counts DATA-table opens /// only and ignores the publisher metadata opens. No-op in production /// (the classification runs only inside the probe closure, which `current` skips -/// when no probes are installed). Called at both open chokepoints. +/// when no probes are installed). Called at the open chokepoint. pub(crate) fn record_open(uri: &str) { let _ = current(|p| { if open_is_internal(uri) { diff --git a/crates/omnigraph/src/table_store.rs b/crates/omnigraph/src/table_store.rs index b149bf2e..e943a3df 100644 --- a/crates/omnigraph/src/table_store.rs +++ b/crates/omnigraph/src/table_store.rs @@ -142,12 +142,18 @@ impl StagedWrite { #[derive(Debug, Clone)] pub struct TableStore { root_uri: String, + /// The graph's shared Lance `Session` (one per graph — LanceDB's + /// one-session-per-connection pattern). Held non-optionally so every open + /// this store performs attaches it; the read path's handle cache shares + /// the same instance via `ReadCaches`. + session: Arc, } impl TableStore { - pub fn new(root_uri: &str) -> Self { + pub fn new(root_uri: &str, session: Arc) -> Self { Self { root_uri: root_uri.trim_end_matches('/').to_string(), + session, } } @@ -194,7 +200,7 @@ impl TableStore { } pub async fn open_at_entry(&self, entry: &SubTableEntry) -> Result { - entry.open(&self.root_uri).await + entry.open(&self.root_uri, Some(&self.session)).await } pub async fn open_dataset_head( @@ -209,7 +215,7 @@ impl TableStore { let ds = crate::instrumentation::open_dataset( dataset_uri, crate::instrumentation::VersionResolution::Latest, - None, + Some(&self.session), crate::instrumentation::table_wrapper(), ) .await?; @@ -226,7 +232,7 @@ impl TableStore { let mut ds = crate::instrumentation::open_dataset( dataset_uri, crate::instrumentation::VersionResolution::Latest, - None, + Some(&self.session), crate::instrumentation::table_wrapper(), ) .await?; @@ -243,7 +249,7 @@ impl TableStore { let ds = crate::instrumentation::open_dataset( dataset_uri, crate::instrumentation::VersionResolution::Latest, - None, + Some(&self.session), crate::instrumentation::table_wrapper(), ) .await?; @@ -271,7 +277,7 @@ impl TableStore { let mut ds = crate::instrumentation::open_dataset( dataset_uri, crate::instrumentation::VersionResolution::Latest, - None, + Some(&self.session), crate::instrumentation::table_wrapper(), ) .await?; diff --git a/crates/omnigraph/tests/failpoints.rs b/crates/omnigraph/tests/failpoints.rs index b3e874f3..9155ab46 100644 --- a/crates/omnigraph/tests/failpoints.rs +++ b/crates/omnigraph/tests/failpoints.rs @@ -1131,7 +1131,7 @@ async fn recovery_rolls_forward_ensure_indices_on_feature_branch() { // publisher deliberately skips the normal index-rebuild preparation; // the failed writer below is still the real `ensure_indices_on`. let person_uri = node_table_uri(&uri, "Person"); - let store = TableStore::new(&uri); + let store = TableStore::new(&uri, std::sync::Arc::new(lance::session::Session::default())); let mut ds = store .open_dataset_head(&person_uri, Some("feature")) .await diff --git a/crates/omnigraph/tests/recovery.rs b/crates/omnigraph/tests/recovery.rs index 4f86f34f..7206b7c8 100644 --- a/crates/omnigraph/tests/recovery.rs +++ b/crates/omnigraph/tests/recovery.rs @@ -1441,7 +1441,7 @@ async fn recovery_classifies_feature_branch_sidecar_against_feature_branch() { // Bypass the manifest: append directly to Person's Lance HEAD on the // feature branch ref to advance HEAD past v_pin. let person_uri = node_table_uri(uri, "Person"); - let store = TableStore::new(uri); + let store = TableStore::new(uri, std::sync::Arc::new(lance::session::Session::default())); let mut ds = store .open_dataset_head(&person_uri, feature_branch_name.as_deref()) .await @@ -1556,7 +1556,7 @@ async fn recovery_rolls_back_feature_branch_sidecar_against_feature_branch() { // Bypass the manifest: append on the feature ref to advance HEAD past // the manifest pin. let person_uri = node_table_uri(uri, "Person"); - let store = TableStore::new(uri); + let store = TableStore::new(uri, std::sync::Arc::new(lance::session::Session::default())); let mut ds = store .open_dataset_head(&person_uri, feature_branch_name.as_deref()) .await diff --git a/crates/omnigraph/tests/staged_writes.rs b/crates/omnigraph/tests/staged_writes.rs index e233e222..e50edfc1 100644 --- a/crates/omnigraph/tests/staged_writes.rs +++ b/crates/omnigraph/tests/staged_writes.rs @@ -123,7 +123,7 @@ fn collect_age_for_id(batches: &[RecordBatch], needle: &str) -> Option { async fn stage_append_is_visible_via_scan_with_staged() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); // Seed: one committed row. let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) @@ -153,7 +153,7 @@ async fn stage_append_is_visible_via_scan_with_staged() { async fn stage_merge_insert_dedupes_superseded_committed_fragment() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); // Seed: alice age 30 in one committed fragment. let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) @@ -215,7 +215,7 @@ async fn stage_merge_insert_dedupes_superseded_committed_fragment() { async fn count_rows_with_staged_matches_scan() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) .await @@ -250,7 +250,7 @@ async fn count_rows_with_staged_matches_scan() { async fn chained_stage_appends_have_distinct_row_ids() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))])) .await @@ -336,7 +336,7 @@ fn combine_for_scan(ds: &Dataset, staged: &[StagedWrite]) -> Vec { async fn stage_append_then_commit_persists_data() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) .await @@ -369,7 +369,7 @@ async fn stage_append_then_commit_persists_data() { async fn stage_merge_insert_then_commit_persists_merged_view() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) .await @@ -401,7 +401,7 @@ async fn stage_merge_insert_then_commit_persists_merged_view() { async fn stage_merge_insert_commit_rebases_over_disjoint_committed_delete() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset(&uri, numbered_person_batch(0..100)) .await @@ -461,7 +461,7 @@ async fn stage_merge_insert_commit_rebases_over_disjoint_committed_delete() { async fn scan_with_staged_with_filter_silently_drops_staged_rows() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); // Committed: alice=30, carol=40 let ds = TableStore::write_dataset( @@ -528,7 +528,7 @@ async fn scan_with_staged_with_filter_silently_drops_staged_rows() { async fn chained_stage_merge_insert_with_shared_key_documents_duplicate_behavior() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); // Seed empty (an unrelated row keeps the schema unambiguous). let ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))])) @@ -594,7 +594,7 @@ async fn chained_stage_merge_insert_with_shared_key_documents_duplicate_behavior async fn stage_overwrite_does_not_advance_head_until_commit() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset(&uri, person_batch(&[("alice", Some(30))])) .await @@ -638,7 +638,7 @@ async fn stage_overwrite_does_not_advance_head_until_commit() { async fn stage_overwrite_preserves_stable_row_ids() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); // `write_dataset` creates with `enable_stable_row_ids: true` — see // ADR 0001. We verify that as a precondition so a future change to @@ -680,7 +680,7 @@ async fn stage_overwrite_preserves_stable_row_ids() { async fn stage_overwrite_replaces_all_fragments() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset( &uri, @@ -719,7 +719,7 @@ async fn stage_overwrite_replaces_all_fragments() { async fn stage_overwrite_empty_batch_replaces_all_rows() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset( &uri, @@ -774,7 +774,7 @@ async fn stage_overwrite_empty_batch_replaces_all_rows() { async fn stage_create_btree_index_does_not_advance_head_until_commit() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset( &uri, @@ -821,7 +821,7 @@ async fn stage_create_btree_index_does_not_advance_head_until_commit() { async fn stage_create_inverted_index_does_not_advance_head_until_commit() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset( &uri, @@ -862,7 +862,7 @@ async fn stage_create_inverted_index_does_not_advance_head_until_commit() { async fn stage_delete_does_not_advance_head_and_reads_through_staged() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset( &uri, @@ -917,7 +917,7 @@ async fn stage_delete_does_not_advance_head_and_reads_through_staged() { async fn stage_delete_commit_rebases_over_disjoint_committed_delete() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let ds = TableStore::write_dataset(&uri, numbered_person_batch(0..100)) .await @@ -955,7 +955,7 @@ async fn create_vector_index_advances_head_inline_documents_residual() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/vec.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); // Build a small dataset with a fixed-size vector column. Vector index // training requires multiple rows; provide enough. @@ -1196,7 +1196,7 @@ async fn commit_staged_skips_auto_cleanup_so_pinned_versions_survive() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let mut ds = TableStore::write_dataset(&uri, person_batch(&[("seed", Some(0))])) .await diff --git a/crates/omnigraph/tests/write_cost.rs b/crates/omnigraph/tests/write_cost.rs index 0b3e1ed9..f58d9d49 100644 --- a/crates/omnigraph/tests/write_cost.rs +++ b/crates/omnigraph/tests/write_cost.rs @@ -92,7 +92,7 @@ async fn internal_table_scans_are_flat_in_history() { /// /// **This is a TRIPWIRE, not the final gate.** It asserts the scan *grows*, i.e. it /// pins the CURRENT served-regime cost (green today) — exactly the `assert_grows` -/// idiom its sibling `data_table_reads_split_into_flat_opener_and_growing_scan` uses, +/// idiom its sibling `data_table_reads_split_into_flat_opener_and_scan_flat_with_session` uses, /// and the "turns red when the fix lands" shape of the Lance surface guards. It flips /// RED the moment the amplification is fixed (write-path probe-gated warm reuse, and /// bringing `__manifest` into `cleanup` version-GC so F stays bounded in history). @@ -141,15 +141,20 @@ async fn internal_table_scans_grow_without_compaction() { // *probe that isolates* the opener (the `PrefixCounter` split) is validated here, // every-PR, on local FS: -/// Proves the `PrefixCounter` opener/scan split: a committing write's data-table -/// reads divide into a **flat opener** term and a **growing scan** term. This pins -/// (a) the classifier actually attributes reads to the opener bucket (non-zero, so a -/// flat assertion isn't vacuously flat-at-zero), and (b) the local data-table growth -/// is the merge-insert/RI fragment scan, not the opener — which is *why* the S3 -/// gate asserts `data_opener_reads`, not total `data_reads`. (On local FS the opener -/// is O(1) regardless of step 3a; the opener's history-dependence is gated on S3.) +/// Proves the `PrefixCounter` opener/scan split — and that BOTH terms are now +/// flat in history on local FS. The opener was always O(1) locally (step 3a); +/// the merge-insert/RI fragment-scan term used to grow with fragment count and +/// was flattened by the dataset-opener unification attaching the shared +/// per-graph `Session` to write-side opens: fragment/manifest metadata is +/// immutable per version, so repeat writes serve it from the session cache +/// instead of re-reading O(depth) objects. This pins (a) the classifier +/// actually attributes reads to the opener bucket (non-zero, so a flat +/// assertion isn't vacuously flat-at-zero), and (b) the session-flattened scan +/// term stays flat — a regression red here means a write-side open dropped the +/// session. (The opener's history-dependence on a real object store stays +/// gated by `write_cost_s3.rs`.) #[tokio::test] -async fn data_table_reads_split_into_flat_opener_and_growing_scan() { +async fn data_table_reads_split_into_flat_opener_and_scan_flat_with_session() { let dir = tempfile::tempdir().unwrap(); let mut db = local_graph(&dir).await; @@ -175,7 +180,10 @@ async fn data_table_reads_split_into_flat_opener_and_growing_scan() { so a flat opener assertion would be vacuous" ); assert_flat(&curve, |c| c.data_opener_reads, 4, "local data-table opener"); - assert_grows(&curve, |c| c.data_scan_reads, 20, "local data-table scan"); + // Pre-session this term GREW with fragment count (the merge-insert/RI scan + // re-reading O(depth) fragment metadata per write); the shared session + // makes repeat reads of immutable metadata cache hits, so it is now flat. + assert_flat(&curve, |c| c.data_scan_reads, 4, "local data-table scan (session-cached)"); } // ── (B) Green-today regression guards — run on every PR ── diff --git a/crates/omnigraph/tests/writes.rs b/crates/omnigraph/tests/writes.rs index 08dedf7c..8e6f4485 100644 --- a/crates/omnigraph/tests/writes.rs +++ b/crates/omnigraph/tests/writes.rs @@ -1411,7 +1411,7 @@ async fn scan_with_pending_rejects_key_column_missing_from_projection() { let dir = tempfile::tempdir().unwrap(); let uri = format!("{}/people.lance", dir.path().to_str().unwrap()); - let store = TableStore::new(dir.path().to_str().unwrap()); + let store = TableStore::new(dir.path().to_str().unwrap(), std::sync::Arc::new(lance::session::Session::default())); let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Utf8, false), diff --git a/docs/dev/testing.md b/docs/dev/testing.md index b53ec548..04ead416 100644 --- a/docs/dev/testing.md +++ b/docs/dev/testing.md @@ -145,7 +145,7 @@ Correctness bugs fail loudly in tests; cost-scaling bugs pass every test and deg - **Assert a cost budget, not just a result.** For a read/open path, assert the number of `Dataset::open` calls (or object-store ops) a warm query performs, and that it does not grow with commit count. The reference is LanceDB's IO-counted tests, which assert a cached read costs 0-1 IO and carry a named regression test against "a list call on every subsequent query." - **Test at history depth.** Build a fixture with many *commits* (not many rows) and assert warm-read cost is flat across depths. A shallow fixture cannot catch an O(commits) cost. -- **Use the shared harness, and gate each term on the backend where it manifests.** `helpers::cost` (`measure`/`IoCounts`/`assert_flat`/`local_graph`/`s3_graph`) is the one place the `IOTracker`/task-local plumbing lives — consume it, don't duplicate it. The write path has *two distinct* depth terms that split cleanly across backends, and conflating them is a real trap (the local data-table read count grows with depth too, but for a different reason — the merge-insert/RI scan reading O(depth) *fragments*, reduced by compaction, not by the opener): (1) the **internal-table** scan term (`__manifest` fragment scans, lineage rows included) reproduces on **any** backend including local FS, so `write_cost.rs` gates it on local every-PR; (2) the **data-table opener** term (latest-version resolution) is a per-object-store-RPC phenomenon — local-FS resolves latest with one cheap `read_dir` regardless of the opener used, so the namespace-vs-direct difference is **invisible on local** and only shows on a real object store (per-version GETs), gated by the bucket-gated `write_cost_s3.rs`. Same harness, different fixture; each term asserted where it actually appears. **`write_cost_s3` is a cost (IO-count) gate, not a correctness test, so it was pulled out of the every-merge `rustfs_integration` CI job — run it on demand (`OMNIGRAPH_S3_TEST_BUCKET=… cargo test -p omnigraph-engine --test write_cost_s3`) pending a dedicated cost/perf harness. The local `write_cost.rs` opener/scan-split guard still runs every-PR, so the split itself stays covered; only the S3 acceptance of the opener term is off the correctness path.** +- **Use the shared harness, and gate each term on the backend where it manifests.** `helpers::cost` (`measure`/`IoCounts`/`assert_flat`/`local_graph`/`s3_graph`) is the one place the `IOTracker`/task-local plumbing lives — consume it, don't duplicate it. The write path has *two distinct* depth terms that split cleanly across backends, and conflating them is a real trap (the local data-table *scan* term used to grow with depth for a different reason — the merge-insert/RI scan re-reading O(depth) *fragments* — until the dataset-opener unification attached the shared per-graph `Session` to write-side opens; immutable fragment/manifest metadata now comes from the session cache, and `write_cost.rs::data_table_reads_split_into_flat_opener_and_scan_flat_with_session` pins that flatness — a red there means a write-side open dropped the session): (1) the **internal-table** scan term (`__manifest` fragment scans, lineage rows included) reproduces on **any** backend including local FS, so `write_cost.rs` gates it on local every-PR; (2) the **data-table opener** term (latest-version resolution) is a per-object-store-RPC phenomenon — local-FS resolves latest with one cheap `read_dir` regardless of the opener used, so the namespace-vs-direct difference is **invisible on local** and only shows on a real object store (per-version GETs), gated by the bucket-gated `write_cost_s3.rs`. Same harness, different fixture; each term asserted where it actually appears. **`write_cost_s3` is a cost (IO-count) gate, not a correctness test, so it was pulled out of the every-merge `rustfs_integration` CI job — run it on demand (`OMNIGRAPH_S3_TEST_BUCKET=… cargo test -p omnigraph-engine --test write_cost_s3`) pending a dedicated cost/perf harness. The local `write_cost.rs` opener/scan-split guard still runs every-PR, so the split itself stays covered; only the S3 acceptance of the opener term is off the correctness path.** - **Count on the handle that does the reads, not just the one a measured op opens.** Lance's IO-counted tests attach the `IOTracker` to the (warm, cached) dataset and read `incremental_stats()` per request — the tracker MUST be on the handle performing the reads, or warm-handle reads escape. A per-op tracker installed at measure time cannot see reads on a long-lived handle opened earlier (the warm coordinator's `__manifest` handle, reused across writes), so such reads were silently undercounted. Wrap a depth-swept body in `cost_harness` so the manifest tracker is installed before the graph opens and `manifest_reads` is **ground truth** (handle-age-irrelevant). The `version_probes` counter is the freshness-probe *call* count; ground truth additionally reveals that a write's probe does ~3 object-store RPCs (a read's probe is a 0-IO cache hit). `manifest_reads_capture_warm_probe` is the guard that this stays true. - This is the testing companion to invariant 15 in [docs/dev/invariants.md](invariants.md) (hot-path cost is bounded by work, not history).