perf(dataset): reuse session-cached manifest on checkout

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(),
@@ -765,6 +765,11 @@ impl Dataset {
         uri: &str,
         session: &Session,
     ) -> Result<Arc<Manifest>> {
+        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,

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;
 
@@ -871,6 +872,76 @@ async fn test_load_manifest_iops() {
     assert_io_eq!(io_stats, read_iops, 1);
 }
 
+#[tokio::test]
+async fn test_checkout_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();
+
+    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,
+                e_tag: None,
+            },
+            Arc::new(dataset.manifest().clone()),
+        )
+        .await;
+    assert!(
+        dataset.checkout_version(version).await.is_err(),
+        "checkout of a version removed from storage must not be served from cache"
+    );
+}
+
 #[rstest]
 #[tokio::test]
 async fn test_write_params(

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

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);
 }