feat: provide timeout parameter for merge_insert (#2378)

Provides the ability to set a timeout for merge insert. The default
underlying timeout is however long the first attempt takes, or if there
are multiple attempts, 30 seconds. This has two use cases:

1. Make the timeout shorter, when you want to fail if it takes too long.
2. Allow taking more time to do retries.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added support for specifying a timeout when performing merge insert
operations in Python, Node.js, and Rust APIs.
- Introduced a new option to control the maximum allowed execution time
for merge inserts, including retry timeout handling.

- **Documentation**
- Updated and added documentation to describe the new timeout option and
its usage in APIs.

- **Tests**
- Added and updated tests to verify correct timeout behavior during
merge insert operations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

python/python/lancedb/merge.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+from datetime import timedelta
 from typing import TYPE_CHECKING, List, Optional
 
 if TYPE_CHECKING:
@@ -31,6 +32,7 @@ class LanceMergeInsertBuilder(object):
         self._when_not_matched_insert_all = False
         self._when_not_matched_by_source_delete = False
         self._when_not_matched_by_source_condition = None
+        self._timeout = None
 
     def when_matched_update_all(
         self, *, where: Optional[str] = None
@@ -81,6 +83,7 @@ class LanceMergeInsertBuilder(object):
         new_data: DATA,
         on_bad_vectors: str = "error",
         fill_value: float = 0.0,
+        timeout: Optional[timedelta] = None,
     ) -> MergeInsertResult:
         """
         Executes the merge insert operation
@@ -98,10 +101,24 @@ class LanceMergeInsertBuilder(object):
             One of "error", "drop", "fill".
         fill_value: float, default 0.
             The value to use when filling vectors. Only used if on_bad_vectors="fill".
+        timeout: Optional[timedelta], default None
+            Maximum time to run the operation before cancelling it.
+
+            By default, there is a 30-second timeout that is only enforced after the
+            first attempt. This is to prevent spending too long retrying to resolve
+            conflicts. For example, if a write attempt takes 20 seconds and fails,
+            the second attempt will be cancelled after 10 seconds, hitting the
+            30-second timeout. However, a write that takes one hour and succeeds on the
+            first attempt will not be cancelled.
+
+            When this is set, the timeout is enforced on all attempts, including
+            the first.
 
         Returns
         -------
         MergeInsertResult
             version: the new version number of the table after doing merge insert.
         """
+        if timeout is not None:
+            self._timeout = timeout
         return self._table._do_merge(self, new_data, on_bad_vectors, fill_value)

python/python/lancedb/table.py

@@ -3716,6 +3716,7 @@ class AsyncTable:
                 when_not_matched_insert_all=merge._when_not_matched_insert_all,
                 when_not_matched_by_source_delete=merge._when_not_matched_by_source_delete,
                 when_not_matched_by_source_condition=merge._when_not_matched_by_source_condition,
+                timeout=merge._timeout,
             ),
         )
 

python/python/tests/test_table.py

@@ -937,7 +937,7 @@ def test_merge_insert(mem_db: DBConnection):
         table.merge_insert("a")
         .when_matched_update_all()
         .when_not_matched_insert_all()
-        .execute(new_data)
+        .execute(new_data, timeout=timedelta(seconds=10))
     )
     assert merge_insert_res.version == 2
     assert merge_insert_res.num_inserted_rows == 1
@@ -1013,6 +1013,12 @@ def test_merge_insert(mem_db: DBConnection):
     expected = pa.table({"a": [2, 4], "b": ["x", "z"]})
     assert table.to_arrow().sort_by("a") == expected
 
+    # timeout
+    with pytest.raises(Exception, match="merge insert timed out"):
+        table.merge_insert("a").when_matched_update_all().execute(
+            new_data, timeout=timedelta(0)
+        )
+
 
 # We vary the data format because there are slight differences in how
 # subschemas are handled in different formats

python/src/table.rs

@@ -649,6 +649,9 @@ impl Table {
             builder
                 .when_not_matched_by_source_delete(parameters.when_not_matched_by_source_condition);
         }
+        if let Some(timeout) = parameters.timeout {
+            builder.timeout(timeout);
+        }
 
         future_into_py(self_.py(), async move {
             let res = builder.execute(Box::new(batches)).await.infer_error()?;
@@ -807,6 +810,7 @@ pub struct MergeInsertParams {
     when_not_matched_insert_all: bool,
     when_not_matched_by_source_delete: bool,
     when_not_matched_by_source_condition: Option<String>,
+    timeout: Option<std::time::Duration>,
 }
 
 #[pyclass]