From 8bdc515f1ec7f1059d401b6af54bc6c96553ad43 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Sat, 10 Oct 2026 04:38:19 +0000 Subject: [PATCH] feat!: use source statistics to plan merge_insert joins `MergeInsertBuilder::execute` now takes any `Scannable`, and the source is routed so Lance's merge join can pick its build side from statistics: * Rescannable sources become a provider reporting their exact row count. * One-shot sources are collected up to `source_collect_threshold_bytes` (default 64MiB). If they end first they are merged from memory; otherwise the prefix is re-chained, spilled for retries, and reported as an inexact lower bound. The threshold can be set per merge or as a connection default. BREAKING CHANGE: `MergeInsertBuilder::execute` and `BaseTable::merge_insert` take a `Scannable` instead of `Box`. Pass a `Box` (e.g. `Box::new(reader) as Box`), a `RecordBatch`, or a `Vec`. Co-Authored-By: Claude Opus 5.5 --- python/src/table.rs | 5 +- rust/lancedb/src/connection.rs | 51 +- rust/lancedb/src/connection/create_table.rs | 13 + rust/lancedb/src/query.rs | 14 +- rust/lancedb/src/remote/table.rs | 84 ++-- rust/lancedb/src/table.rs | 35 +- rust/lancedb/src/table/computed_columns.rs | 15 +- .../src/table/datafusion/scannable_exec.rs | 77 ++- rust/lancedb/src/table/merge.rs | 452 ++++++++++++++++-- rust/lancedb/src/table/merge/lsm.rs | 10 +- rust/lancedb/src/table/refresh.rs | 8 +- 11 files changed, 651 insertions(+), 113 deletions(-) diff --git a/python/src/table.rs b/python/src/table.rs index d82ebbaaf..d8579f6eb 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1601,7 +1601,10 @@ impl Table { } future_into_py(self_.py(), async move { - let res = builder.execute(Box::new(batches)).await.infer_error()?; + let res = builder + .execute(Box::new(batches) as Box) + .await + .infer_error()?; Ok(MergeResult::from(res)) }) } diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 8a0e489c2..6e9fd01af 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -129,6 +129,7 @@ pub struct OpenTableBuilder { parent: Arc, request: OpenTableRequest, embedding_registry: Arc, + merge_insert_source_collect_threshold_bytes: Option, branch: Option, version: Option, } @@ -151,6 +152,7 @@ impl OpenTableBuilder { managed_versioning: None, }, embedding_registry, + merge_insert_source_collect_threshold_bytes: None, branch: None, version: None, } @@ -303,7 +305,10 @@ impl OpenTableBuilder { /// Open the table pub async fn execute(self) -> Result { let table = self.parent.open_table(self.request).await?; - let table = Table::new_with_embedding_registry(table, self.parent, self.embedding_registry); + let table = Table::new_with_embedding_registry(table, self.parent, self.embedding_registry) + .with_merge_insert_source_collect_threshold_bytes( + self.merge_insert_source_collect_threshold_bytes, + ); // "main" is the default branch, so treat it as no branch. let branch = self.branch.filter(|b| b.as_str() != MAIN_BRANCH); match branch { @@ -468,6 +473,7 @@ impl CloneTableBuilder { pub struct Connection { internal: Arc, embedding_registry: Arc, + merge_insert_source_collect_threshold_bytes: Option, } impl std::fmt::Display for Connection { @@ -484,6 +490,7 @@ impl Connection { Self { internal, embedding_registry, + merge_insert_source_collect_threshold_bytes: None, } } @@ -576,6 +583,9 @@ impl Connection { name.into(), initial_data, ) + .merge_insert_source_collect_threshold_bytes( + self.merge_insert_source_collect_threshold_bytes, + ) } /// Create an empty table with a given schema @@ -604,11 +614,14 @@ impl Connection { /// existence. Uncommitted files or a physical `.lance` directory alone do not /// make a table openable. pub fn open_table(&self, name: impl Into) -> OpenTableBuilder { - OpenTableBuilder::new( + let mut builder = OpenTableBuilder::new( self.internal.clone(), name.into(), self.embedding_registry.clone(), - ) + ); + builder.merge_insert_source_collect_threshold_bytes = + self.merge_insert_source_collect_threshold_bytes; + builder } /// Clone a table in the database @@ -1267,6 +1280,7 @@ pub struct ConnectRequest { pub struct ConnectBuilder { request: ConnectRequest, embedding_registry: Option>, + merge_insert_source_collect_threshold_bytes: Option, #[cfg(feature = "remote")] oauth_config: Option, } @@ -1290,6 +1304,7 @@ impl ConnectBuilder { session: None, }, embedding_registry: None, + merge_insert_source_collect_threshold_bytes: None, #[cfg(feature = "remote")] oauth_config: None, } @@ -1412,6 +1427,15 @@ impl ConnectBuilder { self } + /// Default for [`MergeInsertBuilder::source_collect_threshold_bytes`] on + /// tables opened or created through this connection. + /// + /// [`MergeInsertBuilder::source_collect_threshold_bytes`]: crate::table::merge::MergeInsertBuilder::source_collect_threshold_bytes + pub fn merge_insert_source_collect_threshold_bytes(mut self, bytes: usize) -> Self { + self.merge_insert_source_collect_threshold_bytes = Some(bytes); + self + } + /// [`AwsCredential`] to use when connecting to S3. #[cfg(feature = "aws")] #[deprecated(note = "Pass through storage_options instead")] @@ -1606,6 +1630,8 @@ impl ConnectBuilder { embedding_registry: self .embedding_registry .unwrap_or_else(|| Arc::new(MemoryRegistry::new())), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } @@ -1630,6 +1656,8 @@ impl ConnectBuilder { embedding_registry: self .embedding_registry .unwrap_or_else(|| Arc::new(MemoryRegistry::new())), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } else { let internal = Arc::new(ListingDatabase::connect_with_options(&self.request).await?); @@ -1638,6 +1666,8 @@ impl ConnectBuilder { embedding_registry: self .embedding_registry .unwrap_or_else(|| Arc::new(MemoryRegistry::new())), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } } @@ -1675,6 +1705,7 @@ pub struct ConnectNamespaceBuilder { namespace_client_properties: HashMap, read_consistency_interval: Option, embedding_registry: Option>, + merge_insert_source_collect_threshold_bytes: Option, session: Option>, pushdown_operations: HashSet, } @@ -1688,6 +1719,7 @@ impl ConnectNamespaceBuilder { namespace_client_properties: HashMap::new(), read_consistency_interval: None, embedding_registry: None, + merge_insert_source_collect_threshold_bytes: None, session: None, pushdown_operations: HashSet::new(), } @@ -1757,6 +1789,15 @@ impl ConnectNamespaceBuilder { self } + /// Default for [`MergeInsertBuilder::source_collect_threshold_bytes`] on + /// tables opened or created through this connection. + /// + /// [`MergeInsertBuilder::source_collect_threshold_bytes`]: crate::table::merge::MergeInsertBuilder::source_collect_threshold_bytes + pub fn merge_insert_source_collect_threshold_bytes(mut self, bytes: usize) -> Self { + self.merge_insert_source_collect_threshold_bytes = Some(bytes); + self + } + /// Set a custom session for object stores and caching. /// /// By default, a new session with default configuration will be created. @@ -1818,6 +1859,8 @@ impl ConnectNamespaceBuilder { embedding_registry: self .embedding_registry .unwrap_or_else(|| Arc::new(MemoryRegistry::new())), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } } @@ -1850,6 +1893,7 @@ mod test_utils { Self { internal, embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } @@ -1866,6 +1910,7 @@ mod test_utils { Self { internal, embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } } diff --git a/rust/lancedb/src/connection/create_table.rs b/rust/lancedb/src/connection/create_table.rs index 39cc82ec0..1f28a42be 100644 --- a/rust/lancedb/src/connection/create_table.rs +++ b/rust/lancedb/src/connection/create_table.rs @@ -18,6 +18,7 @@ pub struct CreateTableBuilder { parent: Arc, embeddings: Vec<(EmbeddingDefinition, Arc)>, embedding_registry: Arc, + merge_insert_source_collect_threshold_bytes: Option, request: CreateTableRequest, } @@ -32,10 +33,19 @@ impl CreateTableBuilder { parent, embeddings: Vec::new(), embedding_registry, + merge_insert_source_collect_threshold_bytes: None, request: CreateTableRequest::new(name, data), } } + pub(super) fn merge_insert_source_collect_threshold_bytes( + mut self, + bytes: Option, + ) -> Self { + self.merge_insert_source_collect_threshold_bytes = bytes; + self + } + /// Set the mode for creating the table /// /// This controls what happens if a table with the given name already exists @@ -160,6 +170,9 @@ impl CreateTableBuilder { parent.create_table(self.request).await?, parent, embedding_registry, + ) + .with_merge_insert_source_collect_threshold_bytes( + self.merge_insert_source_collect_threshold_bytes, )) } } diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 8fb56ceca..218c58491 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -2358,8 +2358,8 @@ mod tests { datatypes::{Int32Type, UInt8Type}, }; use arrow_array::{ - FixedSizeListArray, Float32Array, Int32Array, RecordBatch, RecordBatchIterator, - StringArray, cast::AsArray, types::Float32Type, + FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray, + types::Float32Type, }; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use datafusion_physical_plan::display::DisplayableExecutionPlan; @@ -2811,10 +2811,7 @@ mod tests { merge.when_not_matched_by_source_delete(Some(r#""PartyAbbrev" = 'D'"#.to_string())); let result = table .base_table() - .merge_insert( - merge, - Box::new(RecordBatchIterator::new(vec![Ok(source)], schema.clone())), - ) + .merge_insert(merge, Box::new(source)) .await .unwrap(); assert_eq!(result.num_deleted_rows, 1); @@ -2830,10 +2827,7 @@ mod tests { .unwrap(); let mut merge = table.merge_insert(&["id"]); merge.when_matched_update_all(Some(r#"target."PartyAbbrev" = 'D'"#.to_string())); - merge - .execute(Box::new(RecordBatchIterator::new(vec![Ok(source)], schema))) - .await - .unwrap(); + merge.execute(source).await.unwrap(); assert_eq!( table .count_rows(Some(r#""PartyAbbrev" = 'U'"#.to_string())) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 119577dc5..82a2d9b9c 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -59,7 +59,7 @@ use crate::{ merge::MergeInsertBuilder, }, }; -use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; +use arrow_array::{LargeBinaryArray, RecordBatch}; use arrow_ipc::reader::{FileReader, StreamReader}; use arrow_schema::{ArrowError, DataType, SchemaRef}; use async_trait::async_trait; @@ -3064,7 +3064,7 @@ impl BaseTable for RemoteTable { async fn merge_insert( &self, mut params: MergeInsertBuilder, - new_data: Box, + mut new_data: Box, ) -> Result { params.canonicalize_filters()?; self.check_mutable().await?; @@ -3079,21 +3079,24 @@ impl BaseTable for RemoteTable { // body stream fails under HTTP2 (issue #2339). The branch, request // timeout header, and merge query params are all applied inside the exec. // - // The public merge_insert API only accepts a `RecordBatchReader`, which - // is not rescannable and so could not be retried directly. To preserve - // the previous retry-on-retryable-status behaviour, buffer the reader - // into memory first: a `Vec` is rescannable, so the outer - // loop can re-execute the plan (and re-stream the body) on each retry. - // This mirrors the old `send_streaming(with_retry=true)` path, which - // likewise buffered the reader to support retries. - let schema = RecordBatchReader::schema(new_data.as_ref()); - let mut batches = new_data.collect::, _>>()?; - // An empty reader still carries a schema. Keep it in an empty batch so - // the buffered source remains scannable and can be replayed on retries. - if batches.is_empty() { - batches.push(RecordBatch::new_empty(schema)); - } - let source: Box = Box::new(batches); + // A one-shot source could not be retried directly. To preserve the + // retry-on-retryable-status behaviour, buffer it into memory first: a + // `Vec` is rescannable, so the outer loop can re-execute + // the plan (and re-stream the body) on each retry. This mirrors the old + // `send_streaming(with_retry=true)` path, which likewise buffered the + // reader to support retries. + let source: Box = if new_data.rescannable() { + new_data + } else { + let schema = new_data.schema(); + let mut batches: Vec = new_data.scan_as_stream().try_collect().await?; + // An empty reader still carries a schema. Keep it in an empty batch so + // the buffered source remains scannable and can be replayed on retries. + if batches.is_empty() { + batches.push(RecordBatch::new_empty(schema)); + } + Box::new(batches) + }; let rescannable = source.rescannable(); let input: Arc = Arc::new(crate::table::datafusion::scannable_exec::ScannableExec::new(source, None)); @@ -3967,6 +3970,8 @@ impl TryFrom for MergeInsertRequest { #[cfg(test)] mod tests { use std::sync::Arc; + + use arrow_array::RecordBatchReader; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use std::{collections::HashMap, pin::Pin}; @@ -4809,7 +4814,11 @@ mod tests { let mut merge = table.merge_insert(&["some_col"]); merge.when_matched_update_all(Some(r#"target."A" > 0"#.into())); - let result = table.base_table().merge_insert(merge, data).await.unwrap(); + let result = table + .base_table() + .merge_insert(merge, Box::new(data)) + .await + .unwrap(); assert_eq!(result.version, if old_server { 0 } else { 43 }); if !old_server { @@ -5017,7 +5026,11 @@ mod tests { let mut merge = table.merge_insert(&["shard_key", "id"]); merge.when_matched_update_all(None); merge.when_not_matched_insert_all(); - let result = table.base_table().merge_insert(merge, data).await.unwrap(); + let result = table + .base_table() + .merge_insert(merge, Box::new(data)) + .await + .unwrap(); assert_eq!(result.num_inserted_rows, 3); } @@ -5041,7 +5054,7 @@ mod tests { let merge = table.merge_insert(&["id", "id"]); let err = table .base_table() - .merge_insert(merge, data) + .merge_insert(merge, Box::new(data)) .await .unwrap_err(); assert!( @@ -5050,20 +5063,32 @@ mod tests { ); } + #[rstest] + #[case::one_shot(false)] + #[case::rescannable(true)] #[tokio::test] - async fn test_merge_insert_retries_on_409() { + async fn test_merge_insert_retries_on_409(#[case] rescannable: bool) { let batch = RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], ) .unwrap(); - let data: Box = Box::new(RecordBatchIterator::new( - [Ok(batch.clone())], - batch.schema(), - )); + // A one-shot reader is buffered so it can be retried; a rescannable + // source is retried directly. + let data: Box = if rescannable { + Box::new(batch) + } else { + let reader: Box = Box::new(RecordBatchIterator::new( + [Ok(batch.clone())], + batch.schema(), + )); + Box::new(reader) + }; + let attempts = Arc::new(AtomicUsize::new(0)); + let handler_attempts = attempts.clone(); // Default parameters - let table = Table::new_with_handler("my_table", |request| { + let table = Table::new_with_handler("my_table", move |request| { assert_eq!(request.method(), "POST"); assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); @@ -5075,15 +5100,18 @@ mod tests { assert!(!params.contains_key("when_matched_update_all_filt")); assert!(!params.contains_key("when_not_matched_by_source_delete_filt")); + handler_attempts.fetch_add(1, Ordering::SeqCst); http::Response::builder().status(409).body("").unwrap() }); + let merge = table.merge_insert(&["some_col"]); let e = table - .merge_insert(&["some_col"]) - .execute(data) + .base_table() + .merge_insert(merge, data) .await .unwrap_err(); assert!(e.to_string().contains("Hit retry limit")); + assert!(attempts.load(Ordering::SeqCst) > 1); } #[rstest] diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index b296d135d..17262ab3d 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -654,7 +654,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { async fn merge_insert( &self, params: MergeInsertBuilder, - new_data: Box, + new_data: Box, ) -> Result; /// Set the unenforced primary key for the table to a single column. /// @@ -998,6 +998,8 @@ pub struct Table { inner: Arc, database: Option>, embedding_registry: Arc, + // Connection default for `MergeInsertBuilder::source_collect_threshold_bytes`. + merge_insert_source_collect_threshold_bytes: Option, } #[cfg(all(test, feature = "remote"))] @@ -1023,6 +1025,7 @@ mod test_utils { database: Some(database), // Registry is unused. embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } @@ -1047,6 +1050,7 @@ mod test_utils { database: Some(database), // Registry is unused. embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } @@ -1069,6 +1073,7 @@ mod test_utils { database: Some(database), // Registry is unused. embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } @@ -1093,6 +1098,7 @@ mod test_utils { database: Some(database), // Registry is unused. embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } @@ -1121,6 +1127,7 @@ mod test_utils { database: Some(database), // Registry is unused. embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } } @@ -1138,6 +1145,7 @@ impl From> for Table { inner, database: None, embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } } @@ -1148,6 +1156,7 @@ impl Table { inner, database: Some(database), embedding_registry: Arc::new(MemoryRegistry::new()), + merge_insert_source_collect_threshold_bytes: None, } } @@ -1177,9 +1186,18 @@ impl Table { inner, database: Some(database), embedding_registry, + merge_insert_source_collect_threshold_bytes: None, } } + pub(crate) fn with_merge_insert_source_collect_threshold_bytes( + mut self, + bytes: Option, + ) -> Self { + self.merge_insert_source_collect_threshold_bytes = bytes; + self + } + /// Cast as [`NativeTable`], or return None it if is not a [`NativeTable`]. /// /// Warning: This function will be removed soon (features exclusive to NativeTable @@ -1222,6 +1240,8 @@ impl Table { inner: self.inner.query_snapshot().await?, database: self.database.clone(), embedding_registry: self.embedding_registry.clone(), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } @@ -1623,13 +1643,14 @@ impl Table { /// merge_insert /// .when_matched_update_all(None) /// .when_not_matched_insert_all(); - /// merge_insert.execute(Box::new(new_data)).await.unwrap(); + /// merge_insert.execute(Box::new(new_data) as Box).await.unwrap(); /// # }); /// ``` pub fn merge_insert(&self, on: &[&str]) -> MergeInsertBuilder { MergeInsertBuilder::new( self.inner.clone(), on.iter().map(|s| s.to_string()).collect(), + self.merge_insert_source_collect_threshold_bytes, ) } @@ -2170,6 +2191,8 @@ impl Table { inner, database: self.database.clone(), embedding_registry: self.embedding_registry.clone(), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } @@ -2469,6 +2492,8 @@ impl Table { inner, database: self.database.clone(), embedding_registry: self.embedding_registry.clone(), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } @@ -2491,6 +2516,8 @@ impl Table { inner, database: self.database.clone(), embedding_registry: self.embedding_registry.clone(), + merge_insert_source_collect_threshold_bytes: self + .merge_insert_source_collect_threshold_bytes, }) } @@ -3613,9 +3640,9 @@ impl BaseTable for NativeTable { async fn merge_insert( &self, params: MergeInsertBuilder, - new_data: Box, + new_data: Box, ) -> Result { - let source_schema = arrow_array::RecordBatchReader::schema(&new_data); + let source_schema = new_data.schema(); self.ensure_not_a_view("merge insert into").await?; computed_columns::ensure_not_written( &Schema::from(self.dataset.get().await?.schema()), diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 88c47992d..6f8c80764 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -3008,13 +3008,7 @@ mod tests { merge .when_matched_update_all(None) .when_not_matched_insert_all(); - let err = merge - .execute(Box::new(arrow_array::RecordBatchIterator::new( - vec![Ok(batch.clone())], - batch.schema(), - ))) - .await - .unwrap_err(); + let err = merge.execute(batch.clone()).await.unwrap_err(); assert!(matches!(err, Error::InvalidInput { .. })); // The append that omits the column still works. @@ -3290,7 +3284,7 @@ mod tests { #[tokio::test] async fn test_unset_with_retained_lsm_rows_cannot_admit_a_declaration() { use crate::table::LsmWriteSpec; - use arrow_array::{Int64Array, RecordBatchIterator}; + use arrow_array::Int64Array; let tmp_dir = tempfile::tempdir().unwrap(); let conn = connect(tmp_dir.path().to_str().unwrap()) @@ -3325,10 +3319,7 @@ mod tests { .when_matched_update_all(None) .when_not_matched_insert_all() .use_lsm(true); - merge - .execute(Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema))) - .await - .unwrap(); + merge.execute(batch).await.unwrap(); table.unset_lsm_write_spec().await.unwrap(); let err = add_computed(&table, &[("doubled".into(), "value * 2".into())]) diff --git a/rust/lancedb/src/table/datafusion/scannable_exec.rs b/rust/lancedb/src/table/datafusion/scannable_exec.rs index b3a4cc536..a5421ac31 100644 --- a/rust/lancedb/src/table/datafusion/scannable_exec.rs +++ b/rust/lancedb/src/table/datafusion/scannable_exec.rs @@ -4,9 +4,16 @@ use core::fmt; use std::sync::{Arc, Mutex}; +use arrow_schema::SchemaRef; +use async_trait::async_trait; +use datafusion_catalog::{Session, TableProvider}; use datafusion_common::{DataFusionError, Result as DFResult, Statistics, stats::Precision}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_expr::{Expr, TableType}; +use datafusion_physical_expr::{ + EquivalenceProperties, Partitioning, PhysicalExpr, expressions::Column, +}; +use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, execution_plan::EmissionType, @@ -118,13 +125,65 @@ impl ExecutionPlan for ScannableExec { } fn partition_statistics(&self, _partition: Option) -> DFResult> { - Ok(Arc::new(Statistics { - num_rows: self - .num_rows - .map(Precision::Exact) - .unwrap_or(Precision::Absent), - total_byte_size: Precision::Absent, - column_statistics: vec![], - })) + // Projections above this node index into `column_statistics`, so it + // needs an (unknown) entry per column. + let mut stats = Statistics::new_unknown(&self.schema()); + if let Some(num_rows) = self.num_rows { + stats.num_rows = Precision::Exact(num_rows); + } + Ok(Arc::new(stats)) + } +} + +/// A [`TableProvider`] over a source plan that can be executed more than once. +/// +/// Every scan re-executes the same plan, so wrapping a [`ScannableExec`] over a +/// rescannable [`Scannable`] yields a provider that replays the source and +/// reports the statistics of the plan. +#[derive(Debug)] +pub(crate) struct ScannableProvider { + plan: Arc, +} + +impl ScannableProvider { + pub fn new(plan: Arc) -> Self { + Self { plan } + } +} + +#[async_trait] +impl TableProvider for ScannableProvider { + fn schema(&self) -> SchemaRef { + self.plan.schema() + } + + fn table_type(&self) -> TableType { + TableType::Temporary + } + + // Filters are never pushed down, and `limit` is only a hint, so the full + // plan is returned with just the projection applied. + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> DFResult> { + let Some(projection) = projection else { + return Ok(self.plan.clone()); + }; + let schema = self.plan.schema(); + let exprs = projection + .iter() + .map(|&idx| { + let name = schema.field(idx).name(); + ( + Arc::new(Column::new(name, idx)) as Arc, + name.clone(), + ) + }) + .collect::>(); + Ok(Arc::new(ProjectionExec::try_new(exprs, self.plan.clone())?)) } } diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 1ad084e22..a72abac53 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -5,16 +5,26 @@ use std::collections::HashSet; use std::sync::Arc; use std::time::Duration; -use arrow_array::RecordBatchReader; +use arrow_array::RecordBatch; use arrow_schema::{DataType, Fields}; +use datafusion::prelude::SessionContext; +use datafusion_catalog::TableProvider; +use datafusion_common::{Statistics, stats::Precision}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_plan::stream::RecordBatchStreamAdapter; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, execute_stream, +}; use futures::future::Either; -use futures::{FutureExt, TryFutureExt}; +use futures::{FutureExt, StreamExt, TryStreamExt}; use lance::dataset::{ MergeInsertBuilder as LanceMergeInsertBuilder, WhenMatched, WhenNotMatchedBySource, }; -use lance_datafusion::utils::StreamingWriteSource; +use lance_datafusion::spill::spilling_table_provider; use serde::{Deserialize, Serialize}; +use super::datafusion::scannable_exec::{ScannableExec, ScannableProvider}; +use crate::data::scannable::Scannable; use crate::error::{Error, Result}; use super::{BaseTable, NativeTable}; @@ -78,10 +88,15 @@ pub struct MergeInsertBuilder { pub(crate) use_index: bool, pub(crate) use_lsm: Option, pub(crate) validate_single_shard: bool, + pub(crate) source_collect_threshold_bytes: Option, } impl MergeInsertBuilder { - pub(super) fn new(table: Arc, on: Vec) -> Self { + pub(super) fn new( + table: Arc, + on: Vec, + source_collect_threshold_bytes: Option, + ) -> Self { Self { table, on, @@ -94,6 +109,7 @@ impl MergeInsertBuilder { use_index: true, use_lsm: None, validate_single_shard: true, + source_collect_threshold_bytes, } } @@ -219,16 +235,44 @@ impl MergeInsertBuilder { self } + /// How many bytes of a one-shot source to collect into memory before + /// streaming the rest. + /// + /// The merge join picks its build side from the statistics each input + /// reports. A rescannable source (such as a [`RecordBatch`] or + /// `Vec`) reports its exact row count. A one-shot source (such + /// as a [`RecordBatchReader`]) reports nothing, so it is read into memory + /// up to this many bytes: if it ends first, it is merged as an in-memory + /// source with exact statistics; otherwise the read rows are reported as a + /// lower bound and the rest is streamed. + /// + /// Overrides the connection default set with + /// [`ConnectBuilder::merge_insert_source_collect_threshold_bytes`](crate::connection::ConnectBuilder::merge_insert_source_collect_threshold_bytes). + /// Defaults to 64MiB. Only applies to local tables. + /// + /// [`RecordBatch`]: arrow_array::RecordBatch + pub fn source_collect_threshold_bytes(&mut self, bytes: usize) -> &mut Self { + self.source_collect_threshold_bytes = Some(bytes); + self + } + /// Executes the merge insert operation /// /// Returns version and statistics about the merge operation including the number of rows /// inserted, updated, and deleted. - pub async fn execute( - mut self, - new_data: Box, - ) -> Result { + /// + /// `new_data` can be any [`Scannable`], such as a [`RecordBatch`], + /// `Vec`, or `Box`. See + /// [`Self::source_collect_threshold_bytes`] for how the source affects the + /// merge plan. + /// + /// [`RecordBatch`]: arrow_array::RecordBatch + pub async fn execute(mut self, new_data: impl Scannable + 'static) -> Result { self.canonicalize_filters()?; - self.table.clone().merge_insert(self, new_data).await + self.table + .clone() + .merge_insert(self, Box::new(new_data)) + .await } pub(crate) fn canonicalize_filters(&mut self) -> Result<()> { @@ -310,7 +354,7 @@ fn validate_merge_source_type(input: &DataType, target: &DataType, path: &str) - pub(crate) async fn execute_merge_insert( table: &NativeTable, mut params: MergeInsertBuilder, - new_data: Box, + new_data: Box, ) -> Result { params.canonicalize_filters()?; super::computed_columns::ensure_no_function_bindings_for_mutation( @@ -338,24 +382,19 @@ pub(crate) async fn execute_merge_insert( let schema = arrow_schema::Schema::from(dataset.schema()); // JSON source fields must carry the stored extension metadata just as on // append. Keep arrow.json text labelled until Lance encodes it as JSONB. - let source = if schema + let cast_schema = if schema .fields() .iter() .any(|field| lance_arrow::json::has_json_fields(field)) { validate_merge_source_fields(new_data.schema().fields(), schema.fields(), "")?; - let plan = Arc::new(super::datafusion::scannable_exec::ScannableExec::new( - Box::new(new_data), - None, - )); - let plan = super::datafusion::cast::cast_to_table_schema(plan, &schema)?; - datafusion_physical_plan::execute_stream( - plan, - Arc::new(datafusion_execution::TaskContext::default()), - )? + Some(&schema) } else { - new_data.into_stream() + None }; + let collect_threshold_bytes = params + .source_collect_threshold_bytes + .unwrap_or(DEFAULT_SOURCE_COLLECT_THRESHOLD_BYTES); let mut builder = LanceMergeInsertBuilder::try_new(dataset.clone(), params.on)?; match ( params.when_matched_update_all, @@ -389,18 +428,28 @@ pub(crate) async fn execute_merge_insert( } builder.use_index(params.use_index); + if let Some(timeout) = params.timeout { + builder.retry_timeout(timeout); + } + let job = builder.try_build()?; + // Reading the source used to happen inside `execute`, so collecting it + // stays under the same timeout. + let future = async move { + match prepare_merge_source(new_data, cast_schema, collect_threshold_bytes).await? { + MergeSource::Provider(provider) => job.execute_provider(provider).await, + MergeSource::Batches(batches) => job.execute_batches(batches).await, + } + .map_err(Error::from) + }; let future = if let Some(timeout) = params.timeout { - let future = builder.retry_timeout(timeout).try_build()?.execute(source); Either::Left(tokio::time::timeout(timeout, future).map(|res| match res { - Ok(Ok((new_dataset, stats))) => Ok((new_dataset, stats)), - Ok(Err(e)) => Err(e.into()), + Ok(result) => result, Err(_) => Err(Error::Runtime { message: "merge insert timed out".to_string(), }), })) } else { - let job = builder.try_build()?; - Either::Right(job.execute(source).map_err(|e| e.into())) + Either::Right(future) }; let (new_dataset, stats) = future.await?; let version = new_dataset.manifest().version; @@ -415,6 +464,142 @@ pub(crate) async fn execute_merge_insert( }) } +/// Default for [`MergeInsertBuilder::source_collect_threshold_bytes`]. +pub const DEFAULT_SOURCE_COLLECT_THRESHOLD_BYTES: usize = 64 * 1024 * 1024; + +// Matches the memory budget Lance uses when it spills a one-shot source itself. +const SPILL_MEMORY_LIMIT_BYTES: usize = 100 * 1024 * 1024; + +/// A merge source in the form Lance plans with. +pub(crate) enum MergeSource { + /// A rescannable provider. Retries rescan it. + Provider(Arc), + /// The whole source, collected into memory. + Batches(Vec), +} + +/// Prepare a merge source so the merge join can see its statistics. +/// +/// * A rescannable source becomes a provider reporting its exact row count. +/// * A one-shot source is read until it ends or exceeds +/// `collect_threshold_bytes`. If it ends, the collected batches are returned. +/// Otherwise the read prefix is chained back onto the rest of the stream, +/// spilled so retries can replay it, and reported with the prefix row count +/// as a lower bound. +/// +/// When `cast_schema` is set, the source is cast to it on every scan. +pub(crate) async fn prepare_merge_source( + source: Box, + cast_schema: Option<&arrow_schema::Schema>, + collect_threshold_bytes: usize, +) -> Result { + let rescannable = source.rescannable(); + let mut plan: Arc = Arc::new(ScannableExec::new(source, None)); + if let Some(schema) = cast_schema { + plan = super::datafusion::cast::cast_to_table_schema(plan, schema)?; + } + if rescannable { + return Ok(MergeSource::Provider(Arc::new(ScannableProvider::new( + plan, + )))); + } + + let schema = plan.schema(); + let mut stream = execute_stream(plan, Arc::new(TaskContext::default()))?; + let mut prefix = Vec::new(); + let mut prefix_bytes = 0; + while prefix_bytes <= collect_threshold_bytes { + match stream.try_next().await? { + Some(batch) => { + prefix_bytes += batch.get_array_memory_size(); + prefix.push(batch); + } + None => { + // An empty batch keeps the source schema, which Lance would + // otherwise replace with the table schema. + if prefix.is_empty() { + prefix.push(RecordBatch::new_empty(schema)); + } + return Ok(MergeSource::Batches(prefix)); + } + } + } + + let prefix_rows = prefix.iter().map(|batch| batch.num_rows()).sum(); + let chained = futures::stream::iter(prefix.into_iter().map(Ok)).chain(stream); + let chained = Box::pin(RecordBatchStreamAdapter::new(schema, chained)); + let spilled = spilling_table_provider(chained, SPILL_MEMORY_LIMIT_BYTES).await?; + let spilled = spilled + .scan(&SessionContext::new().state(), None, &[], None) + .await?; + let plan = Arc::new(RowCountLowerBoundExec::new(spilled, prefix_rows)); + Ok(MergeSource::Provider(Arc::new(ScannableProvider::new( + plan, + )))) +} + +/// Reports an inexact row count for its input, which otherwise reports none. +#[derive(Debug)] +struct RowCountLowerBoundExec { + input: Arc, + num_rows: usize, +} + +impl RowCountLowerBoundExec { + fn new(input: Arc, num_rows: usize) -> Self { + Self { input, num_rows } + } +} + +impl DisplayAs for RowCountLowerBoundExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "RowCountLowerBoundExec: num_rows>={}", self.num_rows) + } +} + +impl ExecutionPlan for RowCountLowerBoundExec { + fn name(&self) -> &str { + "RowCountLowerBoundExec" + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> datafusion_common::Result> { + if children.len() != 1 { + return Err(datafusion_common::DataFusionError::Internal( + "RowCountLowerBoundExec expects exactly one child".to_string(), + )); + } + Ok(Arc::new(Self::new(children.remove(0), self.num_rows))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> datafusion_common::Result { + self.input.execute(partition, context) + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> datafusion_common::Result> { + let mut stats = Statistics::new_unknown(&self.schema()); + stats.num_rows = Precision::Inexact(self.num_rows); + Ok(Arc::new(stats)) + } +} + #[cfg(test)] mod tests { use arrow_array::builder::FixedSizeBinaryBuilder; @@ -425,6 +610,7 @@ mod tests { use arrow_schema::{DataType, Field, Schema}; use std::sync::Arc; + use super::*; use crate::connect; #[rstest::rstest] @@ -501,7 +687,8 @@ mod tests { let mut merge = table.merge_insert(&["id"]); merge.when_not_matched_insert_all(); let error = merge - .execute(Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema))) + .execute(Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)) + as Box) .await .expect_err("JSON alignment must not drop source fields"); let expected = if duplicate { @@ -566,10 +753,10 @@ mod tests { .when_matched_update_all(None) .when_not_matched_insert_all(); let result = merge - .execute(Box::new(RecordBatchIterator::new( - vec![Ok(batch)], - input_schema, - ))) + .execute( + Box::new(RecordBatchIterator::new(vec![Ok(batch)], input_schema)) + as Box, + ) .await .unwrap(); assert_eq!(result.num_updated_rows, 1); @@ -659,6 +846,204 @@ mod tests { Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)) } + fn id_batch(ids: std::ops::Range) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, false), + Field::new("age", DataType::Int32, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(ids.clone())), + Arc::new(Int32Array::from_iter_values(ids)), + ], + ) + .unwrap() + } + + fn one_shot(batches: Vec) -> Box { + let schema = batches[0].schema(); + let reader: Box = Box::new(RecordBatchIterator::new( + batches.into_iter().map(Ok), + schema, + )); + Box::new(reader) + } + + async fn scan_provider(provider: &Arc) -> (Precision, usize) { + let plan = provider + .scan(&SessionContext::new().state(), None, &[], None) + .await + .unwrap(); + let num_rows = plan.partition_statistics(None).unwrap().num_rows; + let batches: Vec = execute_stream(plan, Arc::new(TaskContext::default())) + .unwrap() + .try_collect() + .await + .unwrap(); + (num_rows, batches.iter().map(|b| b.num_rows()).sum()) + } + + #[tokio::test] + async fn prepare_source_rescannable_reports_exact_rows() { + let source: Box = Box::new(vec![id_batch(0..10), id_batch(10..15)]); + let MergeSource::Provider(provider) = prepare_merge_source(source, None, 0).await.unwrap() + else { + panic!("expected a provider"); + }; + // Rescanning replays the whole source each time. + for _ in 0..2 { + assert_eq!(scan_provider(&provider).await, (Precision::Exact(15), 15)); + } + } + + #[tokio::test] + async fn prepare_source_one_shot_under_threshold_collects() { + let source = one_shot(vec![id_batch(0..10), id_batch(10..15)]); + let MergeSource::Batches(batches) = prepare_merge_source(source, None, 1024 * 1024) + .await + .unwrap() + else { + panic!("expected collected batches"); + }; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 15); + } + + #[tokio::test] + async fn prepare_source_one_shot_empty_keeps_schema() { + let empty = RecordBatch::new_empty(id_batch(0..1).schema()); + let source = one_shot(vec![empty.clone()]); + let MergeSource::Batches(batches) = prepare_merge_source(source, None, 1024).await.unwrap() + else { + panic!("expected collected batches"); + }; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema(), empty.schema()); + } + + #[tokio::test] + async fn prepare_source_one_shot_over_threshold_streams_with_lower_bound() { + let first = id_batch(0..10); + let threshold = first.get_array_memory_size(); + let source = one_shot(vec![first, id_batch(10..20), id_batch(20..30)]); + let MergeSource::Provider(provider) = + prepare_merge_source(source, None, threshold).await.unwrap() + else { + panic!("expected a provider"); + }; + // The threshold is exceeded on the second batch, so 20 rows are known. + // The spill lets retries replay the whole source. + for _ in 0..2 { + assert_eq!(scan_provider(&provider).await, (Precision::Inexact(20), 30)); + } + } + + async fn analyze_merge_plan(table: &crate::Table, source: MergeSource) -> String { + let dataset = table.as_native().unwrap().dataset.get().await.unwrap(); + let mut builder = + LanceMergeInsertBuilder::try_new(dataset.clone(), vec!["i".to_string()]).unwrap(); + builder + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(lance::dataset::WhenNotMatched::InsertAll) + .use_index(false); + let job = builder.try_build().unwrap(); + match source { + MergeSource::Provider(provider) => job.analyze_plan_provider(provider).await, + MergeSource::Batches(batches) => job.analyze_plan_batches(batches).await, + } + .unwrap() + } + + /// The input that DataFusion collects into the hash table: the first child + /// of the `HashJoinExec`. + fn hash_join_build_side(plan: &str) -> String { + let lines = plan.lines().collect::>(); + let join = lines + .iter() + .position(|line| line.contains("HashJoinExec")) + .unwrap_or_else(|| panic!("no hash join in plan:\n{plan}")); + let indent = |line: &str| line.len() - line.trim_start().len(); + let child_indent = indent(lines[join]) + 2; + let build_start = join + 1; + let build_end = lines[build_start + 1..] + .iter() + .position(|line| indent(line) <= child_indent) + .map(|offset| build_start + 1 + offset) + .unwrap_or(lines.len()); + lines[build_start..build_end].join("\n") + } + + #[rstest::rstest] + #[case::rescannable(Box::new(id_batch(0..10)) as Box, 1024 * 1024)] + #[case::one_shot_collected(one_shot(vec![id_batch(0..10)]), 1024 * 1024)] + #[tokio::test] + async fn small_source_is_hash_join_build_side( + #[case] source: Box, + #[case] threshold: usize, + ) { + let conn = connect("memory://").execute().await.unwrap(); + let target = (0..10) + .map(|i| id_batch(i * 10_000..(i + 1) * 10_000)) + .collect::>(); + let table = conn.create_table("target", target).execute().await.unwrap(); + + let source = prepare_merge_source(source, None, threshold).await.unwrap(); + let plan = analyze_merge_plan(&table, source).await; + let build_side = hash_join_build_side(&plan); + assert!( + !build_side.contains("LanceRead") && !build_side.contains("LanceScan"), + "target table is the build side:\n{plan}" + ); + } + + #[tokio::test] + async fn merge_insert_one_shot_over_threshold() { + let conn = connect("memory://").execute().await.unwrap(); + let table = conn + .create_table("my_table", id_batch(0..10)) + .execute() + .await + .unwrap(); + + let reader: Box = Box::new(RecordBatchIterator::new( + vec![Ok(id_batch(5..10)), Ok(id_batch(10..20))], + id_batch(0..1).schema(), + )); + let mut merge = table.merge_insert(&["i"]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all() + .source_collect_threshold_bytes(0); + let result = merge.execute(reader).await.unwrap(); + assert_eq!(result.num_updated_rows, 5); + assert_eq!(result.num_inserted_rows, 10); + assert_eq!(table.count_rows(None).await.unwrap(), 20); + } + + #[tokio::test] + async fn merge_insert_threshold_defaults_from_connection() { + let conn = connect("memory://") + .merge_insert_source_collect_threshold_bytes(123) + .execute() + .await + .unwrap(); + let created = conn + .create_table("my_table", id_batch(0..10)) + .execute() + .await + .unwrap(); + let opened = conn.open_table("my_table").execute().await.unwrap(); + for table in [&created, &opened] { + assert_eq!( + table.merge_insert(&["i"]).source_collect_threshold_bytes, + Some(123) + ); + } + let mut merge = opened.merge_insert(&["i"]); + merge.source_collect_threshold_bytes(7); + assert_eq!(merge.source_collect_threshold_bytes, Some(7)); + } + #[tokio::test] async fn test_merge_insert() { let conn = connect("memory://").execute().await.unwrap(); @@ -867,10 +1252,7 @@ mod tests { merge .when_matched_update_all(None) .when_not_matched_by_source_delete(None); - let result = merge - .execute(Box::new(RecordBatchIterator::new([Ok(source)], schema))) - .await - .unwrap(); + let result = merge.execute(source).await.unwrap(); assert_eq!(result.num_updated_rows, 1); assert_eq!(result.num_deleted_rows, (ROW_COUNT - 1) as u64); diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index 0f86fdcac..69e392206 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -24,8 +24,9 @@ use arrow_array::cast::AsArray; use arrow_array::types::{ Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; -use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, RecordBatchReader}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch}; use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef}; +use futures::TryStreamExt; use lance::Dataset; use lance::dataset::mem_wal::{ DatasetMemWalExt, ShardWriter, ShardWriterConfig, evaluate_sharding_spec, @@ -37,6 +38,7 @@ use lance_index::mem_wal::{MemWalIndexDetails, ShardingField, ShardingSpec}; use tokio::sync::RwLock; use uuid::Uuid; +use crate::data::scannable::Scannable; use crate::error::{Error, Result}; use crate::index::IndexConfig; use crate::table::merge::{MergeInsertBuilder, MergeResult}; @@ -693,7 +695,7 @@ pub(crate) async fn execute_lsm_merge_insert( table: &NativeTable, plan: LsmPlan, validate_single_shard: bool, - new_data: Box, + mut new_data: Box, ) -> Result { let dataset = table.dataset.get().await?; let target_schema: SchemaRef = Arc::new(ArrowSchema::from(dataset.schema())); @@ -704,8 +706,8 @@ pub(crate) async fn execute_lsm_merge_insert( let mut batches: Vec = Vec::new(); let mut total_rows: u64 = 0; - for batch in new_data { - let batch = batch.map_err(|e| Error::Arrow { source: e })?; + let mut stream = new_data.scan_as_stream(); + while let Some(batch) = stream.try_next().await? { if batch.num_rows() == 0 { continue; } diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 884648f54..cf0911b10 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -1705,13 +1705,7 @@ mod tests { .when_matched_update_all(None) .when_not_matched_insert_all() .use_lsm(true); - merge - .execute(Box::new(arrow_array::RecordBatchIterator::new( - vec![Ok(batch)], - schema, - ))) - .await - .unwrap(); + merge.execute(batch).await.unwrap(); table.unset_lsm_write_spec().await.unwrap(); super::super::computed_columns::add_foreign_kind(&table, "doubled", "sql").await;