feat: strict overwrite operation (#3712)

Modifies the handling of the Overwrite operation to respect
`read_version` when `num_retries` is set to 0 (strict Overwrite).
A strict Overwrite fails if new commits have been written after
`read_version`.
This allows users to more easily build custom transaction and retry
logic if needed.

The first commit refactors only for code deduplication with no
behavioral changes.

The remainder makes two behavioral changes:
1. Each commit will be tried at least once (previously, commits would
always fail when `num_retries` is set to 0)
2. If the operation is Overwrite and `num_retries` is 0, the commit will
be attempted based on `read_version` instead of attempting to check out
the latest version of the dataset.

Author: HubertY

rust/lance/src/dataset.rs

@@ -3167,6 +3167,41 @@ mod tests {
         assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2);
     }
 
+    #[tokio::test]
+    async fn test_strict_overwrite() {
+        let schema = Schema::try_from(&ArrowSchema::new(vec![ArrowField::new(
+            "x",
+            DataType::Int32,
+            false,
+        )]))
+        .unwrap();
+        let operation = Operation::Overwrite {
+            fragments: vec![],
+            schema,
+            config_upsert_values: None,
+        };
+        let test_dir = tempdir().unwrap();
+        let test_uri = test_dir.path().to_str().unwrap();
+        let read_version_0_transaction = Transaction::new(0, operation, None, None);
+        let strict_builder = CommitBuilder::new(test_uri).with_max_retries(0);
+        let unstrict_builder = CommitBuilder::new(test_uri).with_max_retries(1);
+        strict_builder
+            .clone()
+            .execute(read_version_0_transaction.clone())
+            .await
+            .expect("Strict overwrite should succeed when writing a new dataset");
+        strict_builder
+            .clone()
+            .execute(read_version_0_transaction.clone())
+            .await
+            .expect_err("Strict overwrite should fail when committing to a stale version");
+        unstrict_builder
+            .clone()
+            .execute(read_version_0_transaction.clone())
+            .await
+            .expect("Unstrict overwrite should succeed when committing to a stale version");
+    }
+
     #[rstest]
     #[tokio::test]
     async fn test_merge(

rust/lance/src/io/commit.rs

@@ -732,46 +732,51 @@ pub(crate) async fn commit_transaction(
     // Note: object_store has been configured with WriteParams, but dataset.object_store()
     // has not necessarily. So for anything involving writing, use `object_store`.
     let transaction_file = write_transaction_file(object_store, &dataset.base, transaction).await?;
-
-    // First, get all transactions since read_version
     let read_version = transaction.read_version;
+    let mut target_version = read_version + 1;
     let mut dataset = dataset.clone();
-    // We need to checkout the latest version, because any fixes we apply
-    // (like computing the new row ids) needs to be done based on the most
-    // recent manifest.
-    dataset.checkout_latest().await?;
-    let latest_version = dataset.manifest.version;
-    let other_transactions = futures::stream::iter((read_version + 1)..=latest_version)
-        .map(|version| {
-            read_dataset_transaction_file(&dataset, version)
-                .map(move |res| res.map(|tx| (version, tx)))
-        })
-        .buffer_unordered(dataset.object_store().io_parallelism())
-        .take_while(|res| {
-            futures::future::ready(!matches!(
-                res,
-                Err(crate::Error::NotFound { .. }) | Err(crate::Error::DatasetNotFound { .. })
-            ))
-        })
-        .try_collect::<Vec<_>>()
-        .await?;
-
-    let mut target_version = latest_version + 1;
-
-    if is_detached_version(target_version) {
-        return Err(Error::Internal { message: "more than 2^65 versions have been created and so regular version numbers are appearing as 'detached' versions.".into(), location: location!() });
-    }
-
-    // If any of them conflict with the transaction, return an error
-    for (other_version, other_transaction) in other_transactions.iter() {
-        check_transaction(
-            transaction,
-            *other_version,
-            Some(other_transaction.as_ref()),
-        )?;
+    if matches!(transaction.operation, Operation::Overwrite { .. })
+        && commit_config.num_retries == 0
+    {
+        dataset.checkout_version(transaction.read_version).await?;
+    } else {
+        // We need to checkout the latest version, because any fixes we apply
+        // (like computing the new row ids) needs to be done based on the most
+        // recent manifest.
+        dataset.checkout_latest().await?;
     }
-
-    for attempt_i in 0..commit_config.num_retries {
+    let num_attempts = std::cmp::max(commit_config.num_retries, 1);
+    for attempt_i in 0..num_attempts {
+        // See if we can retry the commit. Try to account for all
+        // transactions that have been committed since the read_version.
+        // Use small amount of backoff to handle transactions that all
+        // started at exact same time better.
+        futures::stream::iter(target_version..=dataset.manifest.version)
+            .map(|version| {
+                read_dataset_transaction_file(&dataset, version)
+                    .map(move |res| res.map(|tx| (version, tx)))
+            })
+            .buffer_unordered(dataset.object_store().io_parallelism())
+            .take_while(|res| {
+                futures::future::ready(
+                    attempt_i > 0
+                        || !matches!(
+                            res,
+                            Err(crate::Error::NotFound { .. })
+                                | Err(crate::Error::DatasetNotFound { .. })
+                        ),
+                )
+            })
+            .try_for_each(|(other_version, other_transaction)| {
+                let res =
+                    check_transaction(transaction, other_version, Some(other_transaction.as_ref()));
+                futures::future::ready(res)
+            })
+            .await?;
+        target_version = dataset.manifest.version + 1;
+        if is_detached_version(target_version) {
+            return Err(Error::Internal { message: "more than 2^65 versions have been created and so regular version numbers are appearing as 'detached' versions.".into(), location: location!() });
+        }
         // Build an up-to-date manifest from the transaction and current manifest
         let (mut manifest, mut indices) = match transaction.operation {
             Operation::Restore { version } => {
@@ -843,32 +848,11 @@ pub(crate) async fn commit_transaction(
                 return Ok((manifest, manifest_location.path, manifest_location.e_tag));
             }
             Err(CommitError::CommitConflict) => {
-                // See if we can retry the commit. Try to account for all
-                // transactions that have been committed since the read_version.
-                // Use small amount of backoff to handle transactions that all
-                // started at exact same time better.
-
-                let backoff_time = backoff_time(attempt_i);
-                tokio::time::sleep(backoff_time).await;
-
-                dataset.checkout_latest().await?;
-                let latest_version = dataset.manifest.version;
-                futures::stream::iter(target_version..=latest_version)
-                    .map(|version| {
-                        read_dataset_transaction_file(&dataset, version)
-                            .map(move |res| res.map(|tx| (version, tx)))
-                    })
-                    .buffer_unordered(dataset.object_store().io_parallelism())
-                    .try_for_each(|(version, other_transaction)| {
-                        let res = check_transaction(
-                            transaction,
-                            version,
-                            Some(other_transaction.as_ref()),
-                        );
-                        futures::future::ready(res)
-                    })
-                    .await?;
-                target_version = latest_version + 1;
+                let next_attempt_i = attempt_i + 1;
+                if next_attempt_i < num_attempts {
+                    tokio::time::sleep(backoff_time(next_attempt_i)).await;
+                    dataset.checkout_latest().await?;
+                }
             }
             Err(CommitError::OtherError(err)) => {
                 // If other error, return