From fa7638d48d523cebbda905034eb816901509fce6 Mon Sep 17 00:00:00 2001 From: zhangyue19921010 Date: Tue, 7 Jul 2026 05:37:21 -0400 Subject: [PATCH 1/6] perf(dataset): reuse session-cached manifest on checkout --- rust/lance/src/dataset.rs | 4 +- rust/lance/src/dataset/tests/dataset_io.rs | 65 +++++++++++++++++++++- 2 files changed, 66 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 470c6873dc7..d8ddeb979be 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -599,7 +599,7 @@ impl Dataset { return Ok(self.clone()); } - let manifest = Self::load_manifest( + let manifest = Self::get_manifest( self.object_store.as_ref(), &manifest_location, &new_location.uri, @@ -625,7 +625,7 @@ impl Dataset { self.object_store.clone(), new_location.path, new_location.uri, - Arc::new(manifest), + manifest, manifest_location, self.session.clone(), self.commit_handler.clone(), diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 1f8c7226bf2..d5e6a7fd45a 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -37,7 +37,7 @@ use lance_file::{ version::LanceFileVersion, writer::{FileWriter, FileWriterOptions}, }; -use lance_io::assert_io_eq; +use lance_io::{assert_io_eq, assert_io_gt}; use lance_table::feature_flags; use lance_table::format::BasePath; use object_store::ObjectStoreExt; @@ -871,6 +871,69 @@ async fn test_load_manifest_iops() { assert_io_eq!(io_stats, read_iops, 1); } +#[tokio::test] +async fn test_checkout_reuses_cached_manifest() { + let test_uri = TempStrDir::default(); + let session = Arc::new(Session::default()); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let make_batches = || { + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], + ) + .unwrap(); + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()) + }; + + // v1: committing on this Session caches the v1 manifest. + Dataset::write( + make_batches(), + &test_uri, + Some(WriteParams { + session: Some(session.clone()), + ..Default::default() + }), + ) + .await + .unwrap(); + + // v2 (append): the dataset now points at v2, but the v1 manifest stays in + // the Session cache from the write above. + let dataset = Dataset::write( + make_batches(), + &test_uri, + Some(WriteParams { + session: Some(session.clone()), + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.manifest().version, 2); + + // Cache hit: checking out v1 resolves the manifest location (1 head) and + // reuses the cached manifest body, so no manifest read is issued. + let _ = dataset.object_store.as_ref().io_stats_incremental(); // reset + let ds_v1 = dataset.checkout_version(1u64).await.unwrap(); + assert_eq!(ds_v1.manifest().version, 1); + let io_stats = dataset.object_store.as_ref().io_stats_incremental(); + assert_io_eq!(io_stats, read_iops, 1); + + // Cache miss: clearing the metadata cache forces the same checkout to read + // the manifest body from storage again, costing more than the lone head. + session.file_metadata_cache().clear().await; + let _ = dataset.object_store.as_ref().io_stats_incremental(); // reset + let ds_v1_cold = dataset.checkout_version(1u64).await.unwrap(); + assert_eq!(ds_v1_cold.manifest().version, 1); + let io_stats = dataset.object_store.as_ref().io_stats_incremental(); + assert_io_gt!(io_stats, read_iops, 1); +} + #[rstest] #[tokio::test] async fn test_write_params( From 57c53d979017bbeaabafe3fbccc69d49cddce5cf Mon Sep 17 00:00:00 2001 From: zhangyue19921010 Date: Tue, 7 Jul 2026 07:47:01 -0400 Subject: [PATCH 2/6] perf(dataset): reuse session-cached manifest on checkout --- rust/lance/src/dataset/write/commit.rs | 5 ++++- rust/lance/src/io/commit/s3_test.rs | 7 +++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index d76c2049873..f8d09d3c55e 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -629,8 +629,11 @@ mod tests { assert_eq!(new_ds.manifest().version, 7); // Session should still be re-used // However, the dataset needs to be loaded and the read version checked out. + // The read version's manifest body is served from the session cache (it + // was cached when v1 was first created), so the checkout only pays the + // version-resolution head, not a manifest read. let io_stats = dataset.object_store.as_ref().io_stats_incremental(); - assert_io_eq!(io_stats, read_iops, 4, "load dataset + check version"); + assert_io_eq!(io_stats, read_iops, 3, "load dataset + check version"); assert_io_eq!(io_stats, write_iops, 2, "write txn + manifest"); // Commit transaction with URI and new session. Re-use the store diff --git a/rust/lance/src/io/commit/s3_test.rs b/rust/lance/src/io/commit/s3_test.rs index b5b1a09c776..4be469ee368 100644 --- a/rust/lance/src/io/commit/s3_test.rs +++ b/rust/lance/src/io/commit/s3_test.rs @@ -341,7 +341,10 @@ async fn test_ddb_open_iops() { // Checkout original version dataset.checkout_version(1).await.unwrap(); let io_stats = dataset.object_store.as_ref().io_stats_incremental(); - // Checkout: 1 IOPS: manifest file - assert_io_eq!(io_stats, read_iops, 1); + // Checkout: 0 read IOPS. Version 1's manifest was already loaded and cached + // on this Session when the dataset was opened above, so the checkout serves + // the manifest body from the metadata cache. Version resolution is handled + // in DynamoDB and issues no S3 read. + assert_io_eq!(io_stats, read_iops, 0); assert_io_eq!(io_stats, write_iops, 0); } From 925487f5fdd8fd2972a0f2c51f65eb51e5dd3a7f Mon Sep 17 00:00:00 2001 From: zhangyue19921010 Date: Tue, 7 Jul 2026 09:31:54 -0400 Subject: [PATCH 3/6] perf(dataset): reuse session-cached manifest on checkout --- rust/lance/src/dataset.rs | 5 + rust/lance/src/dataset/tests/dataset_io.rs | 126 +++++++++++++++++++++ 2 files changed, 131 insertions(+) diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index d8ddeb979be..77eaa72f7fe 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -765,6 +765,11 @@ impl Dataset { uri: &str, session: &Session, ) -> Result> { + if manifest_location.size.is_none() { + return Ok(Arc::new( + Self::load_manifest(object_store, manifest_location, uri, session).await?, + )); + } let metadata_cache = session.metadata_cache.for_dataset(uri); let manifest_key = ManifestKey { version: manifest_location.version, diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index d5e6a7fd45a..4672e4c4066 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -12,6 +12,7 @@ use crate::dataset::WriteMode::Overwrite; use crate::dataset::builder::DatasetBuilder; use crate::dataset::{ManifestWriteConfig, write_manifest_file}; use crate::session::Session; +use crate::session::caches::ManifestKey; use crate::{Dataset, Error, Result}; use lance_table::format::DataStorageFormat; @@ -934,6 +935,131 @@ async fn test_checkout_reuses_cached_manifest() { assert_io_gt!(io_stats, read_iops, 1); } +#[tokio::test] +async fn test_checkout_removed_version_not_served_from_cache() { + // Regression: a version that no longer exists in storage (e.g. removed by + // auto-cleanup) must never be served from the session metadata cache. + // Version resolution falls back to an unchecked location (no `size`) for a + // missing version, and the cache key collapses to `manifest/{version}` when + // the store's head yields no e_tag, so a stale cached entry would otherwise + // be returned instead of surfacing a NotFound. + let test_uri = TempStrDir::default(); + let session = Arc::new(Session::default()); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], + ) + .unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_uri, + Some(WriteParams { + session: Some(session.clone()), + ..Default::default() + }), + ) + .await + .unwrap(); + + // Seed a zombie manifest for a version absent from storage, keyed without an + // e_tag to mimic a store whose head yields none (the key collapses to + // `manifest/999`, exactly what a missing version resolves to). Give it the + // matching version so it stands in for a version that was cached and then + // removed (the realistic auto-cleanup scenario). + let mut zombie = dataset.manifest().clone(); + zombie.version = 999; + session + .metadata_cache + .for_dataset(&dataset.uri) + .insert_with_key( + &ManifestKey { + version: 999, + e_tag: None, + }, + Arc::new(zombie), + ) + .await; + + // The checkout must fail rather than return the cached zombie manifest. + assert!( + dataset.checkout_version(999u64).await.is_err(), + "checkout of a version absent from storage must not be served from cache" + ); +} + +#[tokio::test] +async fn test_open_removed_version_not_served_from_cache() { + // Regression: the same zombie-cache hazard applies to opening a specific + // version by URI (`DatasetBuilder::with_version`), not just checkout — both + // resolve a missing version to an unchecked location and go through + // `get_manifest`. + let test_uri = TempStrDir::default(); + let session = Arc::new(Session::default()); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], + ) + .unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_uri, + Some(WriteParams { + session: Some(session.clone()), + ..Default::default() + }), + ) + .await + .unwrap(); + + // Open once through the builder to learn the exact URI it caches under + // (`from_uri` re-normalizes the input, so it may differ from the URI the + // writer recorded). + let opened = DatasetBuilder::from_uri(dataset.uri.as_str()) + .with_session(session.clone()) + .load() + .await + .unwrap(); + + // Seed a zombie manifest for a version absent from storage under that exact + // URI, keyed without an e_tag. Give it the matching version so it passes the + // builder's version-consistency check and would be returned if the cache + // were trusted — i.e. a version that was cached and then removed. + let mut zombie = opened.manifest().clone(); + zombie.version = 999; + session + .metadata_cache + .for_dataset(&opened.uri) + .insert_with_key( + &ManifestKey { + version: 999, + e_tag: None, + }, + Arc::new(zombie), + ) + .await; + + // Opening the absent version by URI must fail, not return the cached zombie. + let result = DatasetBuilder::from_uri(dataset.uri.as_str()) + .with_version(999u64) + .with_session(session.clone()) + .load() + .await; + assert!( + result.is_err(), + "opening a version absent from storage must not be served from cache" + ); +} + #[rstest] #[tokio::test] async fn test_write_params( From 0262022ca92240000fc467a9d6c44b3d975c3891 Mon Sep 17 00:00:00 2001 From: zhangyue19921010 Date: Tue, 7 Jul 2026 09:33:05 -0400 Subject: [PATCH 4/6] perf(dataset): reuse session-cached manifest on checkout --- rust/lance/src/dataset/tests/dataset_io.rs | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 4672e4c4066..0fea2cd6d0e 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -937,12 +937,6 @@ async fn test_checkout_reuses_cached_manifest() { #[tokio::test] async fn test_checkout_removed_version_not_served_from_cache() { - // Regression: a version that no longer exists in storage (e.g. removed by - // auto-cleanup) must never be served from the session metadata cache. - // Version resolution falls back to an unchecked location (no `size`) for a - // missing version, and the cache key collapses to `manifest/{version}` when - // the store's head yields no e_tag, so a stale cached entry would otherwise - // be returned instead of surfacing a NotFound. let test_uri = TempStrDir::default(); let session = Arc::new(Session::default()); let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( @@ -966,11 +960,6 @@ async fn test_checkout_removed_version_not_served_from_cache() { .await .unwrap(); - // Seed a zombie manifest for a version absent from storage, keyed without an - // e_tag to mimic a store whose head yields none (the key collapses to - // `manifest/999`, exactly what a missing version resolves to). Give it the - // matching version so it stands in for a version that was cached and then - // removed (the realistic auto-cleanup scenario). let mut zombie = dataset.manifest().clone(); zombie.version = 999; session @@ -994,10 +983,6 @@ async fn test_checkout_removed_version_not_served_from_cache() { #[tokio::test] async fn test_open_removed_version_not_served_from_cache() { - // Regression: the same zombie-cache hazard applies to opening a specific - // version by URI (`DatasetBuilder::with_version`), not just checkout — both - // resolve a missing version to an unchecked location and go through - // `get_manifest`. let test_uri = TempStrDir::default(); let session = Arc::new(Session::default()); let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( @@ -1030,10 +1015,6 @@ async fn test_open_removed_version_not_served_from_cache() { .await .unwrap(); - // Seed a zombie manifest for a version absent from storage under that exact - // URI, keyed without an e_tag. Give it the matching version so it passes the - // builder's version-consistency check and would be returned if the cache - // were trusted — i.e. a version that was cached and then removed. let mut zombie = opened.manifest().clone(); zombie.version = 999; session From 9e41dda35e40a95d412632af1e1f217d5d8f7ed8 Mon Sep 17 00:00:00 2001 From: zhangyue19921010 Date: Tue, 7 Jul 2026 09:38:06 -0400 Subject: [PATCH 5/6] perf(dataset): reuse session-cached manifest on checkout --- rust/lance/src/dataset/tests/dataset_io.rs | 128 +-------------------- 1 file changed, 3 insertions(+), 125 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 0fea2cd6d0e..f221cad0271 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -38,7 +38,7 @@ use lance_file::{ version::LanceFileVersion, writer::{FileWriter, FileWriterOptions}, }; -use lance_io::{assert_io_eq, assert_io_gt}; +use lance_io::assert_io_eq; use lance_table::feature_flags; use lance_table::format::BasePath; use object_store::ObjectStoreExt; @@ -872,69 +872,6 @@ async fn test_load_manifest_iops() { assert_io_eq!(io_stats, read_iops, 1); } -#[tokio::test] -async fn test_checkout_reuses_cached_manifest() { - let test_uri = TempStrDir::default(); - let session = Arc::new(Session::default()); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let make_batches = || { - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], - ) - .unwrap(); - RecordBatchIterator::new(vec![Ok(batch)], schema.clone()) - }; - - // v1: committing on this Session caches the v1 manifest. - Dataset::write( - make_batches(), - &test_uri, - Some(WriteParams { - session: Some(session.clone()), - ..Default::default() - }), - ) - .await - .unwrap(); - - // v2 (append): the dataset now points at v2, but the v1 manifest stays in - // the Session cache from the write above. - let dataset = Dataset::write( - make_batches(), - &test_uri, - Some(WriteParams { - session: Some(session.clone()), - mode: WriteMode::Append, - ..Default::default() - }), - ) - .await - .unwrap(); - assert_eq!(dataset.manifest().version, 2); - - // Cache hit: checking out v1 resolves the manifest location (1 head) and - // reuses the cached manifest body, so no manifest read is issued. - let _ = dataset.object_store.as_ref().io_stats_incremental(); // reset - let ds_v1 = dataset.checkout_version(1u64).await.unwrap(); - assert_eq!(ds_v1.manifest().version, 1); - let io_stats = dataset.object_store.as_ref().io_stats_incremental(); - assert_io_eq!(io_stats, read_iops, 1); - - // Cache miss: clearing the metadata cache forces the same checkout to read - // the manifest body from storage again, costing more than the lone head. - session.file_metadata_cache().clear().await; - let _ = dataset.object_store.as_ref().io_stats_incremental(); // reset - let ds_v1_cold = dataset.checkout_version(1u64).await.unwrap(); - assert_eq!(ds_v1_cold.manifest().version, 1); - let io_stats = dataset.object_store.as_ref().io_stats_incremental(); - assert_io_gt!(io_stats, read_iops, 1); -} - #[tokio::test] async fn test_checkout_removed_version_not_served_from_cache() { let test_uri = TempStrDir::default(); @@ -960,6 +897,8 @@ async fn test_checkout_removed_version_not_served_from_cache() { .await .unwrap(); + // A cached manifest for a version removed from storage (keyed without an + // e_tag) must not be served: the checkout has to fail, not return the zombie. let mut zombie = dataset.manifest().clone(); zombie.version = 999; session @@ -974,73 +913,12 @@ async fn test_checkout_removed_version_not_served_from_cache() { ) .await; - // The checkout must fail rather than return the cached zombie manifest. assert!( dataset.checkout_version(999u64).await.is_err(), "checkout of a version absent from storage must not be served from cache" ); } -#[tokio::test] -async fn test_open_removed_version_not_served_from_cache() { - let test_uri = TempStrDir::default(); - let session = Arc::new(Session::default()); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], - ) - .unwrap(); - let dataset = Dataset::write( - RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), - &test_uri, - Some(WriteParams { - session: Some(session.clone()), - ..Default::default() - }), - ) - .await - .unwrap(); - - // Open once through the builder to learn the exact URI it caches under - // (`from_uri` re-normalizes the input, so it may differ from the URI the - // writer recorded). - let opened = DatasetBuilder::from_uri(dataset.uri.as_str()) - .with_session(session.clone()) - .load() - .await - .unwrap(); - - let mut zombie = opened.manifest().clone(); - zombie.version = 999; - session - .metadata_cache - .for_dataset(&opened.uri) - .insert_with_key( - &ManifestKey { - version: 999, - e_tag: None, - }, - Arc::new(zombie), - ) - .await; - - // Opening the absent version by URI must fail, not return the cached zombie. - let result = DatasetBuilder::from_uri(dataset.uri.as_str()) - .with_version(999u64) - .with_session(session.clone()) - .load() - .await; - assert!( - result.is_err(), - "opening a version absent from storage must not be served from cache" - ); -} - #[rstest] #[tokio::test] async fn test_write_params( From fb5248a4f8f979334a74b5e3b9f7ef439fe5c6a0 Mon Sep 17 00:00:00 2001 From: zhangyue19921010 Date: Thu, 9 Jul 2026 03:19:12 -0400 Subject: [PATCH 6/6] code review --- rust/lance/src/dataset/tests/dataset_io.rs | 47 ++++++++++++++++------ 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index f221cad0271..f9618914037 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -897,25 +897,48 @@ async fn test_checkout_removed_version_not_served_from_cache() { .await .unwrap(); - // A cached manifest for a version removed from storage (keyed without an - // e_tag) must not be served: the checkout has to fail, not return the zombie. - let mut zombie = dataset.manifest().clone(); - zombie.version = 999; - session - .metadata_cache - .for_dataset(&dataset.uri) + let version = dataset.manifest().version; + let location = dataset.manifest_location().clone(); + let cache = session.metadata_cache.for_dataset(&dataset.uri); + + assert!( + cache + .get_with_key(&ManifestKey { + version, + e_tag: location.e_tag.as_deref(), + }) + .await + .is_some(), + "manifest should be cached after the write" + ); + dataset.checkout_version(version).await.unwrap(); + + // Remove the version from storage, as cleanup (or a manual delete) would. + dataset.object_store.delete(&location.path).await.unwrap(); + + let resolved = dataset + .commit_handler + .resolve_version_location(&dataset.base, version, &dataset.object_store.inner) + .await + .unwrap(); + assert!( + resolved.size.is_none(), + "resolving a removed version must fall back to a size-less location, got {:?}", + resolved.size + ); + + cache .insert_with_key( &ManifestKey { - version: 999, + version, e_tag: None, }, - Arc::new(zombie), + Arc::new(dataset.manifest().clone()), ) .await; - assert!( - dataset.checkout_version(999u64).await.is_err(), - "checkout of a version absent from storage must not be served from cache" + dataset.checkout_version(version).await.is_err(), + "checkout of a version removed from storage must not be served from cache" ); }