From 1c229a90ea86f09ffcef53240c445b052103c501 Mon Sep 17 00:00:00 2001
From: albertlockett <albert@lancedb.com>
Date: Thu, 18 Jul 2024 15:09:58 -0300
Subject: [PATCH] support passing WriteParams to update

---
 Cargo.toml                | 15 ++++++++++-----
 rust/lancedb/src/table.rs | 14 ++++++++++++++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index eafa1cdd..10104d9b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,11 +20,16 @@ keywords = ["lancedb", "lance", "database", "vector", "search"]
 categories = ["database-implementations"]
 
 [workspace.dependencies]
-lance = { "version" = "=0.14.2", "features" = ["dynamodb"] }
-lance-index = { "version" = "=0.14.2" }
-lance-linalg = { "version" = "=0.14.2" }
-lance-testing = { "version" = "=0.14.2" }
-lance-datafusion = { "version" = "=0.14.2" }
+# lance = { "version" = "=0.14.2", "features" = ["dynamodb"] }
+# lance-index = { "version" = "=0.14.2" }
+# lance-linalg = { "version" = "=0.14.2" }
+# lance-testing = { "version" = "=0.14.2" }
+# lance-datafusion = { "version" = "=0.14.2" }
+lance = { path = "../lance/rust/lance", "features" = ["dynamodb"] }
+lance-index = { path = "../lance/rust/lance-index" }
+lance-linalg = { path = "../lance/rust/lance-linalg" }
+lance-testing = { path = "../lance/rust/lance-testing" }
+lance-datafusion = { path = "../lance/rust/lance-datafusion" }
 # Note that this one does not include pyarrow
 arrow = { version = "51.0", optional = false }
 arrow-array = "51.0"
diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs
index 4aa698d3..f477eb0b 100644
--- a/rust/lancedb/src/table.rs
+++ b/rust/lancedb/src/table.rs
@@ -298,6 +298,7 @@ pub struct UpdateBuilder {
     parent: Arc<dyn TableInternal>,
     pub(crate) filter: Option<String>,
     pub(crate) columns: Vec<(String, String)>,
+    pub(crate) write_options: Option<WriteOptions>
 }
 
 impl UpdateBuilder {
@@ -306,6 +307,7 @@ impl UpdateBuilder {
             parent,
             filter: None,
             columns: Vec::new(),
+            write_options: None,
         }
     }
 
@@ -345,6 +347,12 @@ impl UpdateBuilder {
         self
     }
 
+    /// Apply the given write options when updating the dataset
+    pub fn write_options(mut self, write_options: Option<WriteOptions>) -> Self {
+        self.write_options = write_options;
+        self
+    }
+
     /// Executes the update operation
     pub async fn execute(self) -> Result<()> {
         if self.columns.is_empty() {
@@ -1673,6 +1681,12 @@ impl TableInternal for NativeTable {
             builder = builder.set(column, &value)?;
         }
 
+        if let Some(write_options) = update.write_options {
+            if let Some(write_params) = write_options.lance_write_params {
+                builder = builder.with_write_params(write_params);
+            }
+        }
+
         let operation = builder.build()?;
         let ds = operation.execute().await?;
         self.dataset.set_latest(ds.as_ref().clone()).await;