rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-12 14:52:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat!: use source statistics to plan merge_insert joins

Browse Source 
`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<dyn RecordBatchReader + Send>`. Pass a `Box<dyn RecordBatchReader +
Send>` (e.g. `Box::new(reader) as Box<dyn RecordBatchReader + Send>`), a
`RecordBatch`, or a `Vec<RecordBatch>`.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Will Jones
2026-10-10 04:38:19 +00:00
co-authored by Claude Opus 5.5
parent c2c8269fc8
commit 8bdc515f1e
11 changed files with 651 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/table.rs
+4 -1
View File
@@ -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<dyn arrow::array::RecordBatchReader + Send>)
.await
.infer_error()?;
Ok(MergeResult::from(res))
})
}
rust/lancedb/src/connection.rs
+48 -3
View File
@@ -129,6 +129,7 @@ pub struct OpenTableBuilder {
parent: Arc<dyn Database>,
request: OpenTableRequest,
embedding_registry: Arc<dyn EmbeddingRegistry>,
merge_insert_source_collect_threshold_bytes: Option<usize>,
branch: Option<String>,
version: Option<u64>,
}
@@ -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<Table> {
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<dyn Database>,
embedding_registry: Arc<dyn EmbeddingRegistry>,
merge_insert_source_collect_threshold_bytes: Option<usize>,
}
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 `<name>.lance` directory alone do not
/// make a table openable.
pub fn open_table(&self, name: impl Into<String>) -> 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<Arc<dyn EmbeddingRegistry>>,
merge_insert_source_collect_threshold_bytes: Option<usize>,
#[cfg(feature = "remote")]
oauth_config: Option<crate::remote::OAuthConfig>,
}
@@ -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<String, String>,
read_consistency_interval: Option<std::time::Duration>,
embedding_registry: Option<Arc<dyn EmbeddingRegistry>>,
merge_insert_source_collect_threshold_bytes: Option<usize>,
session: Option<Arc<lance::session::Session>>,
pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
}
@@ -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,
}
}
}
rust/lancedb/src/connection/create_table.rs
+13
View File
@@ -18,6 +18,7 @@ pub struct CreateTableBuilder {
parent: Arc<dyn Database>,
embeddings: Vec<(EmbeddingDefinition, Arc<dyn EmbeddingFunction>)>,
embedding_registry: Arc<dyn EmbeddingRegistry>,
merge_insert_source_collect_threshold_bytes: Option<usize>,
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<usize>,
) -> 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,
))
}
}
rust/lancedb/src/query.rs
+4 -10
View File
@@ -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()))
rust/lancedb/src/remote/table.rs
+56 -28
View File
@@ -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<S: HttpSend> BaseTable for RemoteTable<S> {
async fn merge_insert(
&self,
mut params: MergeInsertBuilder,
new_data: Box<dyn RecordBatchReader + Send>,
mut new_data: Box<dyn Scannable>,
) -> Result<MergeResult> {
params.canonicalize_filters()?;
self.check_mutable().await?;
@@ -3079,21 +3079,24 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
// 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<RecordBatch>` 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::<std::result::Result<Vec<_>, _>>()?;
// 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<dyn Scannable> = 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<RecordBatch>` 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<dyn Scannable> = if new_data.rescannable() {
new_data
} else {
let schema = new_data.schema();
let mut batches: Vec<RecordBatch> = 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<dyn ExecutionPlan> =
Arc::new(crate::table::datafusion::scannable_exec::ScannableExec::new(source, None));
@@ -3967,6 +3970,8 @@ impl TryFrom<MergeInsertBuilder> 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<dyn RecordBatchReader + Send> = 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<dyn Scannable> = if rescannable {
Box::new(batch)
} else {
let reader: Box<dyn RecordBatchReader + Send> = 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]
rust/lancedb/src/table.rs
+31 -4
View File
@@ -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<dyn RecordBatchReader + Send>,
new_data: Box<dyn Scannable>,
) -> Result<MergeResult>;
/// Set the unenforced primary key for the table to a single column.
///
@@ -998,6 +998,8 @@ pub struct Table {
inner: Arc<dyn BaseTable>,
database: Option<Arc<dyn Database>>,
embedding_registry: Arc<dyn EmbeddingRegistry>,
// Connection default for `MergeInsertBuilder::source_collect_threshold_bytes`.
merge_insert_source_collect_threshold_bytes: Option<usize>,
}
#[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<Arc<dyn BaseTable>> 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<usize>,
) -> 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<dyn RecordBatchReader + Send>).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<dyn RecordBatchReader + Send>,
new_data: Box<dyn Scannable>,
) -> Result<MergeResult> {
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()),
rust/lancedb/src/table/computed_columns.rs
+3 -12
View File
@@ -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())])
rust/lancedb/src/table/datafusion/scannable_exec.rs
+68 -9
View File
@@ -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<usize>) -> DFResult<Arc<Statistics>> {
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<dyn ExecutionPlan>,
}
impl ScannableProvider {
pub fn new(plan: Arc<dyn ExecutionPlan>) -> 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<usize>>,
_filters: &[Expr],
_limit: Option<usize>,
) -> DFResult<Arc<dyn ExecutionPlan>> {
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<dyn PhysicalExpr>,
name.clone(),
)
})
.collect::<Vec<_>>();
Ok(Arc::new(ProjectionExec::try_new(exprs, self.plan.clone())?))
}
}
rust/lancedb/src/table/merge.rs
+417 -35
View File
@@ -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<bool>,
pub(crate) validate_single_shard: bool,
pub(crate) source_collect_threshold_bytes: Option<usize>,
}
impl MergeInsertBuilder {
pub(super) fn new(table: Arc<dyn BaseTable>, on: Vec<String>) -> Self {
pub(super) fn new(
table: Arc<dyn BaseTable>,
on: Vec<String>,
source_collect_threshold_bytes: Option<usize>,
) -> 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<RecordBatch>`) 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<dyn RecordBatchReader + Send>,
) -> Result<MergeResult> {
///
/// `new_data` can be any [`Scannable`], such as a [`RecordBatch`],
/// `Vec<RecordBatch>`, or `Box<dyn RecordBatchReader + Send>`. 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<MergeResult> {
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<dyn RecordBatchReader + Send>,
new_data: Box<dyn Scannable>,
) -> Result<MergeResult> {
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<dyn TableProvider>),
/// The whole source, collected into memory.
Batches(Vec<RecordBatch>),
}
/// 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<dyn Scannable>,
cast_schema: Option<&arrow_schema::Schema>,
collect_threshold_bytes: usize,
) -> Result<MergeSource> {
let rescannable = source.rescannable();
let mut plan: Arc<dyn ExecutionPlan> = 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<dyn ExecutionPlan>,
num_rows: usize,
}
impl RowCountLowerBoundExec {
fn new(input: Arc<dyn ExecutionPlan>, 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<PlanProperties> {
self.input.properties()
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![&self.input]
}
fn with_new_children(
self: Arc<Self>,
mut children: Vec<Arc<dyn ExecutionPlan>>,
) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
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<TaskContext>,
) -> datafusion_common::Result<SendableRecordBatchStream> {
self.input.execute(partition, context)
}
fn partition_statistics(
&self,
_partition: Option<usize>,
) -> datafusion_common::Result<Arc<Statistics>> {
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<dyn RecordBatchReader + Send>)
.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<dyn RecordBatchReader + Send>,
)
.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<i32>) -> 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<RecordBatch>) -> Box<dyn Scannable> {
let schema = batches[0].schema();
let reader: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
batches.into_iter().map(Ok),
schema,
));
Box::new(reader)
}
async fn scan_provider(provider: &Arc<dyn TableProvider>) -> (Precision<usize>, 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<RecordBatch> = 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<dyn Scannable> = 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::<usize>(), 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::<Vec<_>>();
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<dyn Scannable>, 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<dyn Scannable>,
#[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::<Vec<_>>();
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<dyn RecordBatchReader + Send> = 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);
rust/lancedb/src/table/merge/lsm.rs
+6 -4
View File
@@ -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<dyn RecordBatchReader + Send>,
mut new_data: Box<dyn Scannable>,
) -> Result<MergeResult> {
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<RecordBatch> = 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;
}
rust/lancedb/src/table/refresh.rs
+1 -7
View File
@@ -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;