diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 470c6873dc7..77eaa72f7fe 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(), @@ -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 1f8c7226bf2..f9618914037 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; @@ -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( 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); }