support passing WriteParams to update

This commit is contained in:
albertlockett
2024-07-18 15:09:58 -03:00
parent cc114ada72
commit 1c229a90ea
2 changed files with 24 additions and 5 deletions

View File

@@ -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"

View File

@@ -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;