diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index ac72d566192..70e96e323ab 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -230,6 +230,23 @@ impl From<&Manifest> for Version { } } +/// The transaction that produced a version of the dataset, along with the +/// version's commit timestamp. +/// +/// Returned by [`Dataset::read_version_transaction`], which reads this +/// information directly from storage without checking out the version. +#[derive(Debug, Clone)] +pub struct VersionTransaction { + /// Version number. + pub version: u64, + + /// Timestamp the version was committed, in UTC. + pub timestamp: DateTime, + + /// The transaction that produced this version, if one was recorded. + pub transaction: Option, +} + /// Customize read behavior of a dataset. #[derive(Clone, Debug)] pub struct ReadParams { @@ -1187,43 +1204,146 @@ impl Dataset { return Ok(Some((*transaction).clone())); } + let transaction = self + .read_transaction_from_storage(&self.manifest, &self.manifest_location) + .await?; + + if let Some(tx) = transaction.as_ref() { + self.metadata_cache + .insert_with_key(&transaction_key, Arc::new(tx.clone())) + .await; + } + Ok(transaction) + } + + /// Read the transaction recorded by `manifest` directly from storage, + /// without consulting or populating any session cache. + async fn read_transaction_from_storage( + &self, + manifest: &Manifest, + manifest_location: &ManifestLocation, + ) -> Result> { // Prefer inline transaction from manifest when available - let transaction = if let Some(pos) = self.manifest.transaction_section { - let reader = if let Some(size) = self.manifest_location.size { - self.object_store - .open_with_size(&self.manifest_location.path, size as usize) - .await? - } else { - self.object_store.open(&self.manifest_location.path).await? + if let Some(pos) = manifest.transaction_section { + let reader = match manifest_location.size { + Some(size) => { + self.object_store + .open_with_size(&manifest_location.path, size as usize) + .await? + } + None => self.object_store.open(&manifest_location.path).await?, }; - let tx: pb::Transaction = read_message(reader.as_ref(), pos).await?; - Transaction::try_from(tx).map(Some)? - } else if let Some(path) = &self.manifest.transaction_file { + // A concurrent overwrite can leave the listed size too small; retry + // once with the true size. + let tx: pb::Transaction = match read_message(reader.as_ref(), pos).await { + Err(e) + if manifest_location.size.is_some() + && e.to_string().contains("file size is too small") => + { + let reader = self.object_store.open(&manifest_location.path).await?; + read_message(reader.as_ref(), pos).await? + } + other => other?, + }; + Transaction::try_from(tx).map(Some) + } else if let Some(path) = &manifest.transaction_file { // Fallback: read external transaction file if present let path = self.transactions_dir().join(path.as_str()); let data = self.object_store.inner.get(&path).await?.bytes().await?; let transaction = lance_table::format::pb::Transaction::decode(data)?; - Transaction::try_from(transaction).map(Some)? + Transaction::try_from(transaction).map(Some) } else { - None - }; + Ok(None) + } + } - if let Some(tx) = transaction.as_ref() { - self.metadata_cache - .insert_with_key(&transaction_key, Arc::new(tx.clone())) - .await; + /// Read the transaction (if any) and commit timestamp of a version of the + /// dataset. `version` is a version number on this dataset's current branch. + /// + /// Reads the version's manifest transiently: no historical `Dataset` is + /// constructed, no `IndexSection` is decoded, and no session cache is read + /// or written, so scanning many historical versions does not fill the + /// shared caches. + /// + /// Returns an error if the version does not exist (for example, if it has + /// been cleaned up). + /// + /// # Example + /// + /// ``` + /// # use lance::{Dataset, Result}; + /// # async fn example(dataset: &Dataset) -> Result<()> { + /// let record = dataset.read_version_transaction(5).await?; + /// let committed_at = record.timestamp; + /// let operation = record.transaction.as_ref().map(|t| t.operation.name()); + /// # Ok(()) + /// # } + /// ``` + pub async fn read_version_transaction(&self, version: u64) -> Result { + // Resolve against this dataset's current branch. + let manifest_location = self + .commit_handler + .resolve_version_location(&self.base, version, &self.object_store.inner) + .await?; + + // Keep the DatasetNotFound variant callers expect for a missing version. + let manifest = read_manifest( + &self.object_store, + &manifest_location.path, + manifest_location.size, + ) + .await + .map_err(|e| match &e { + Error::NotFound { uri, .. } => Error::dataset_not_found(uri.clone(), box_error(e)), + _ => e, + })?; + + // The resolved manifest must belong to this dataset's branch. A + // mismatch means the commit handler resolved against a different chain + // (for example an external manifest store that ignores + // branch-qualified paths); error loudly rather than hand back another + // branch's transaction. + if manifest.branch != self.manifest.branch { + return Err(Error::internal(format!( + "reading version {} on branch '{}' resolved a manifest belonging to branch '{}'", + version, + refs::normalize_branch(self.manifest.branch.as_deref()), + refs::normalize_branch(manifest.branch.as_deref()), + ))); } - Ok(transaction) + + let transaction = self + .read_transaction_from_storage(&manifest, &manifest_location) + .await?; + + Ok(VersionTransaction { + version: manifest.version, + timestamp: manifest.timestamp(), + transaction, + }) } /// Read the transaction file for this version of the dataset. /// /// If there was no transaction file written for this version of the dataset /// then this will return None. + /// + /// Does not populate the session caches; see + /// [`Self::read_version_transaction`]. + /// + /// # Example + /// + /// ``` + /// # use lance::{Dataset, Result}; + /// # async fn example(dataset: &Dataset) -> Result<()> { + /// let transaction = dataset.read_transaction_by_version(5).await?; + /// let operation = transaction.as_ref().map(|t| t.operation.name()); + /// # Ok(()) + /// # } + /// ``` pub async fn read_transaction_by_version(&self, version: u64) -> Result> { - let dataset_version = self.checkout_version(version).await?; - dataset_version.read_transaction().await + Ok(self.read_version_transaction(version).await?.transaction) } /// List transactions for the dataset, up to a maximum number. diff --git a/rust/lance/src/dataset/tests/dataset_transactions.rs b/rust/lance/src/dataset/tests/dataset_transactions.rs index 3e2a4caa3b3..74790d301c1 100644 --- a/rust/lance/src/dataset/tests/dataset_transactions.rs +++ b/rust/lance/src/dataset/tests/dataset_transactions.rs @@ -272,6 +272,39 @@ pub(super) fn assert_results( ) } +fn gen_rows() -> impl arrow_array::RecordBatchReader + Send + 'static { + lance_datagen::gen_batch() + .col("key", array::step::()) + .into_reader_rows(RowCount::from(10), BatchCount::from(1)) +} + +/// Write a dataset with `versions` versions of 10 rows each. +async fn write_versions(uri: &str, versions: usize, enable_v2_manifest_paths: bool) -> Dataset { + let mut ds = Dataset::write( + gen_rows(), + uri, + Some(WriteParams { + enable_v2_manifest_paths, + ..Default::default() + }), + ) + .await + .unwrap(); + for _ in 1..versions { + ds.append( + gen_rows(), + Some(WriteParams { + mode: WriteMode::Append, + enable_v2_manifest_paths, + ..Default::default() + }), + ) + .await + .unwrap(); + } + ds +} + #[tokio::test] async fn test_inline_transaction() { use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator}; @@ -382,6 +415,197 @@ async fn test_inline_transaction() { assert!(ds_new.manifest.transaction_file.is_some()); let read_tx = ds_new.read_transaction().await.unwrap().unwrap(); assert_eq!(read_tx, tx); + + // The direct read takes the same external-file fallback. + let version_transaction = ds_new + .read_version_transaction(location.version) + .await + .unwrap(); + assert_eq!(version_transaction.transaction, Some(tx)); +} + +#[tokio::test] +async fn test_read_version_transaction_does_not_populate_caches() { + use lance_index::IndexType; + use lance_index::scalar::ScalarIndexParams; + + let test_uri = TempStrDir::default(); + let mut dataset = write_versions(&test_uri, 1, true).await; + // Index the table so historical manifests carry an IndexSection that a + // caching read path would decode. + dataset + .create_index( + &["key"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); // version 2 + for _ in 0..18 { + dataset + .append( + gen_rows(), + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + } + let latest_version = dataset.version().version; + assert_eq!(latest_version, 20); + + // Fresh session so any cache insertion by the API under test shows as growth. + let session = Arc::new(Session::default()); + let dataset = DatasetBuilder::from_uri(&test_uri) + .with_session(session.clone()) + .load() + .await + .unwrap(); + + let metadata_stats_before = session.metadata_cache_stats().await; + let index_stats_before = session.index_cache_stats().await; + + let mut actual = Vec::with_capacity(latest_version as usize); + for version in 1..=latest_version { + let version_transaction = dataset.read_version_transaction(version).await.unwrap(); + assert_eq!(version_transaction.version, version); + actual.push(version_transaction); + } + + let metadata_stats_after = session.metadata_cache_stats().await; + let index_stats_after = session.index_cache_stats().await; + assert_eq!( + metadata_stats_after.num_entries, + metadata_stats_before.num_entries + ); + assert_eq!( + metadata_stats_after.size_bytes, + metadata_stats_before.size_bytes + ); + assert_eq!( + index_stats_after.num_entries, + index_stats_before.num_entries + ); + assert_eq!(index_stats_after.size_bytes, index_stats_before.size_bytes); + + // Results match a full checkout. + for version_transaction in &actual { + let checked_out = dataset + .checkout_version(version_transaction.version) + .await + .unwrap(); + assert_eq!( + version_transaction.transaction, + checked_out.read_transaction().await.unwrap() + ); + assert_eq!( + version_transaction.timestamp, + checked_out.version().timestamp + ); + assert!(version_transaction.transaction.is_some()); + } + + // A missing (e.g. cleaned up) version errors as DatasetNotFound, matching + // the historical checkout_version-based contract of the public API. + let err = dataset.read_version_transaction(9999).await.unwrap_err(); + assert!( + matches!(err, crate::Error::DatasetNotFound { .. }), + "expected DatasetNotFound for a missing version, got {err:?}" + ); +} + +#[tokio::test] +async fn test_read_transaction_recovers_from_stale_manifest_size() { + let test_uri = TempStrDir::default(); + let ds = write_versions(&test_uri, 1, true).await; + let manifest = ds.manifest().clone(); + // Only meaningful for the inline path; a plain write inlines the transaction. + assert!(manifest.transaction_section.is_some()); + + // A size at/under the transaction offset makes the first read_message fail + // "file size is too small"; only the retry at the true size can recover. + let mut stale = ds.manifest_location().clone(); + stale.size = Some(1); + let recovered = ds + .read_transaction_from_storage(&manifest, &stale) + .await + .unwrap(); + assert_eq!(recovered, ds.read_transaction().await.unwrap()); + assert!(recovered.is_some()); +} + +#[tokio::test] +async fn test_read_version_transaction_v1_manifest_naming() { + let test_uri = TempStrDir::default(); + let ds = write_versions(&test_uri, 3, false).await; + assert_eq!( + ds.manifest_location().naming_scheme, + ManifestNamingScheme::V1 + ); + + for version in 1..=3 { + let version_transaction = ds.read_version_transaction(version).await.unwrap(); + let checked_out = ds.checkout_version(version).await.unwrap(); + assert_eq!( + version_transaction.transaction, + checked_out.read_transaction().await.unwrap() + ); + assert_eq!( + version_transaction.timestamp, + checked_out.version().timestamp + ); + } +} + +#[tokio::test] +async fn test_read_version_transaction_on_branch() { + let test_uri = TempStrDir::default(); + let mut main_ds = write_versions(&test_uri, 1, true).await; + let branch_ds = main_ds.create_branch("dev", 1, None).await.unwrap(); + + // Commit on the branch. + let branch_ds = Dataset::write( + gen_rows(), + branch_ds.uri(), + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(branch_ds.manifest().branch.as_deref(), Some("dev")); + + // Versions resolve against the branch chain and match a full checkout. + for version in branch_ds.versions().await.unwrap() { + let version_transaction = branch_ds + .read_version_transaction(version.version) + .await + .unwrap(); + assert_eq!(version_transaction.version, version.version); + assert_eq!(version_transaction.timestamp, version.timestamp); + let checked_out = branch_ds.checkout_version(version.version).await.unwrap(); + assert_eq!(checked_out.manifest().branch.as_deref(), Some("dev")); + assert_eq!( + version_transaction.transaction, + checked_out.read_transaction().await.unwrap() + ); + } + + // The append on the branch is the branch's own transaction. + let latest = branch_ds.version().version; + let version_transaction = branch_ds.read_version_transaction(latest).await.unwrap(); + assert!(matches!( + version_transaction.transaction, + Some(Transaction { + operation: Operation::Append { .. }, + .. + }) + )); } #[tokio::test]