This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new a0bf047d Support Paimon table sources in prepared MERGE (#1071)
a0bf047d is described below
commit a0bf047d44132f13d1ec75be81b3ba3d26a70c74
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Oct 8 12:00:07 2026 +0800
Support Paimon table sources in prepared MERGE (#1071)
---
bindings/python/README.md | 6 +
.../python/python/pypaimon_rust/datafusion.pyi | 14 +-
bindings/python/src/write.rs | 9 +-
bindings/python/tests/test_update.py | 42 ++++
crates/paimon/src/table/mod.rs | 1 +
crates/paimon/src/table/table_merge_input.rs | 174 ++++++++++++++++
crates/paimon/src/table/table_merge_into.rs | 147 ++++++++------
crates/paimon/src/table/table_merge_into_tests.rs | 223 ++++++++++++++++++++-
8 files changed, 553 insertions(+), 63 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index 8cdca711..3effeb7f 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -123,6 +123,12 @@ Self-merge uses the target's pinned row IDs. A later
action failure retains file
when the operation cannot return their messages. Returned messages may only be
explicitly aborted by a caller that knows publication has been abandoned.
+The source may also be a native `Table`. Core resolves its branch and snapshot
+selectors once, projects the source keys and referenced columns, and reads the
+logical rows through the ordinary reader, including primary-key merge engines.
+The source does not require data evolution or row tracking. Arrow batches keep
+their chunk boundaries; matching gathers only the contributing source chunks.
+
### Custom Blob URI sources
Before the first Arrow write, call
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 8e91f096..52b26b71 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -263,10 +263,13 @@ class BatchTableWrite:
class BatchTableUpdate:
def merge_into(
- self, source: Optional[pyarrow.Table], *, on: Sequence[Tuple[str,
str]],
+ self, source: Optional[Union[pyarrow.Table, Table]], *, on:
Sequence[Tuple[str, str]],
when_matched: Sequence[Dict[str, Any]], when_not_matched:
Sequence[Dict[str, Any]]
) -> List[CommitMessage]:
- """Prepare MERGE against one snapshot. None selects self-merge on
_ROW_ID.
+ """Prepare MERGE against one target snapshot. None selects self-merge
on _ROW_ID.
+
+ A Table source reads its selected full snapshot in core, independently
+ of the target. Arrow input preserves batches and source row order.
Clauses carry assignments as (column, source/target/literal, value)
tuples, a delete flag, and optional SQL condition metadata.
@@ -301,11 +304,14 @@ class TableUpdateByRowId:
class StreamTableUpdate:
def merge_into(
- self, source: Optional[pyarrow.Table], *, on: Sequence[Tuple[str,
str]],
+ self, source: Optional[Union[pyarrow.Table, Table]], *, on:
Sequence[Tuple[str, str]],
when_matched: Sequence[Dict[str, Any]], when_not_matched:
Sequence[Dict[str, Any]],
commit_identifier: int
) -> List[CommitMessage]:
- """Prepare MERGE against one snapshot. None selects self-merge on
_ROW_ID.
+ """Prepare MERGE against one target snapshot. None selects self-merge
on _ROW_ID.
+
+ A Table source reads its selected full snapshot in core, independently
+ of the target. Arrow input preserves batches and source row order.
Clauses carry assignments as (column, source/target/literal, value)
tuples, a delete flag, and optional SQL condition metadata.
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index 7555cebc..ea65b1be 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -592,11 +592,18 @@ impl UpdateContext {
when_not_matched: &Bound<'_, PyAny>,
) -> PyResult<Vec<PyCommitMessage>> {
let source = match source {
- Some(source) =>
paimon::table::MergeSource::Batches(arrow_table_batches(source)?),
+ Some(source) => match
source.extract::<PyRef<crate::table::PyTable>>() {
+ Ok(table) =>
paimon::table::MergeSource::Table(table.inner.clone()),
+ Err(_) =>
paimon::table::MergeSource::Batches(arrow_table_batches(source)?),
+ },
None => paimon::table::MergeSource::SelfTable,
};
let source_schema = match &source {
paimon::table::MergeSource::Batches(batches) =>
batches[0].schema(),
+ paimon::table::MergeSource::Table(table) => {
+
paimon::arrow::build_target_arrow_schema(table.schema().fields())
+ .map_err(to_py_err)?
+ }
paimon::table::MergeSource::SelfTable => {
let schema =
paimon::arrow::build_target_arrow_schema(self.table.schema().fields())
.map_err(to_py_err)?;
diff --git a/bindings/python/tests/test_update.py
b/bindings/python/tests/test_update.py
index 7bed73e2..68157654 100644
--- a/bindings/python/tests/test_update.py
+++ b/bindings/python/tests/test_update.py
@@ -466,3 +466,45 @@ def
test_prepared_merge_conditions_chunks_and_lazy_literals(tmp_path, stream):
commit = builder.new_commit()
commit.commit(42, messages) if stream else commit.commit(messages)
assert rows() == [dict(id=1, value=100), dict(id=2, value=200)]
+
+
[email protected]('stream', [False, True])
[email protected]('historical', [False, True])
+def test_prepared_merge_accepts_native_table_sources(tmp_path, stream,
historical):
+ import json
+ from pathlib import Path
+
+ context = SQLContext()
+ context.register_catalog('paimon', {'warehouse': str(tmp_path)})
+ context.sql('CREATE SCHEMA paimon.table_merge')
+ context.sql("""CREATE TABLE paimon.table_merge.target (id INT, value INT)
WITH (
+ 'row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true')""")
+ # An ordinary append source needs neither row tracking nor data evolution.
+ context.sql('CREATE TABLE paimon.table_merge.source (source_id INT, amount
INT, flag INT)')
+ context.sql('INSERT INTO paimon.table_merge.target (id, value) VALUES (1,
10), (2, 20)')
+ for values in ['(1, 11, 1)', '(3, 33, 1)', '(2, 22, 1)']:
+ context.sql('INSERT INTO paimon.table_merge.source (source_id, amount,
flag) VALUES ' + values)
+ catalog = PaimonCatalog({'warehouse': str(tmp_path)})
+ target, source = catalog.get_table('table_merge.target'),
catalog.get_table('table_merge.source')
+ if historical:
+ schema = json.loads((Path(source.location()) / 'schema' /
'schema-0').read_text())
+ schema['options']['scan.snapshot-id'] = '2'
+ source = source.copy_with_resolved_schema(json.dumps(schema))
+ source_snapshot = source.latest_snapshot().id()
+ target_snapshot = target.latest_snapshot().id()
+ builder =
(target.new_stream_write_builder().with_commit_user('source-merge')
+ if stream else target.new_batch_write_builder())
+ kwargs = dict(commit_identifier=42) if stream else {}
+ condition = dict(sql='"s.flag" = 1')
+ messages = builder.new_update().merge_into(
+ source, on=[('id', 'source_id')], when_matched=[
+ dict(delete=False, condition=condition, assignments=[('value',
'source', 'amount')])],
+ when_not_matched=[dict(condition=condition, assignments=[
+ ('id', 'source', 'source_id'), ('value', 'source', 'amount')])],
**kwargs)
+ assert target.latest_snapshot().id() == target_snapshot
+ assert source.latest_snapshot().id() == source_snapshot
+ commit = builder.new_commit()
+ commit.commit(42, messages) if stream else commit.commit(messages)
+ actual = pa.Table.from_batches(context.sql(
+ 'SELECT id, value FROM paimon.table_merge.target ORDER BY
id')).to_pylist()
+ assert actual == [dict(id=1, value=11), dict(id=2, value=20 if historical
else 22), dict(id=3, value=33)]
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 24a3cf22..6200bc6b 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -142,6 +142,7 @@ mod sorted_global_index_options;
mod source;
mod stats_filter;
pub(crate) mod table_commit;
+mod table_merge_input;
mod table_merge_into;
#[cfg(test)]
mod table_merge_into_tests;
diff --git a/crates/paimon/src/table/table_merge_input.rs
b/crates/paimon/src/table/table_merge_input.rs
new file mode 100644
index 00000000..a0a03d4c
--- /dev/null
+++ b/crates/paimon/src/table/table_merge_input.rs
@@ -0,0 +1,174 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Materialized MERGE input preserves Arrow batches and source row order.
+
+use super::Table;
+use arrow_array::RecordBatch;
+use arrow_schema::SchemaRef;
+use arrow_select::interleave::interleave_record_batch;
+use futures::TryStreamExt;
+use std::collections::{HashMap, HashSet};
+
+fn invalid(message: impl Into<String>) -> crate::Error {
+ crate::Error::DataInvalid {
+ message: message.into(),
+ source: None,
+ }
+}
+
+pub(super) struct MergeInput {
+ pub schema: SchemaRef,
+ pub batches: Vec<RecordBatch>,
+ ends: Vec<usize>,
+ pub row_count: usize,
+}
+
+impl MergeInput {
+ pub fn new(mut batches: Vec<RecordBatch>) -> crate::Result<Self> {
+ let first = batches
+ .first()
+ .ok_or_else(|| invalid("MERGE source needs an Arrow schema"))?;
+ let schema = first.schema();
+ if batches.iter().any(|batch| batch.schema() != schema) {
+ return Err(invalid("MERGE source batches must share a schema"));
+ }
+ let names = schema
+ .fields()
+ .iter()
+ .map(|field| field.name())
+ .collect::<HashSet<_>>();
+ if names.len() != schema.fields().len() {
+ return Err(invalid("MERGE source has duplicate column names"));
+ }
+ batches.retain(|batch| batch.num_rows() > 0);
+ if batches.is_empty() {
+ batches.push(RecordBatch::new_empty(schema.clone()));
+ }
+ let mut row_count = 0usize;
+ let ends = batches
+ .iter()
+ .map(|batch| {
+ row_count = row_count
+ .checked_add(batch.num_rows())
+ .ok_or_else(|| invalid("MERGE source row count
overflow"))?;
+ Ok(row_count)
+ })
+ .collect::<crate::Result<_>>()?;
+ Ok(Self {
+ schema,
+ batches,
+ ends,
+ row_count,
+ })
+ }
+
+ pub async fn read(table: &Table, columns: &HashSet<String>) ->
crate::Result<Self> {
+ if table.is_format_table() {
+ return Err(crate::Error::Unsupported {
+ message: "MERGE table sources require a Paimon table".into(),
+ });
+ }
+
crate::spec::CoreOptions::new(table.schema().options()).validate_scan_options()?;
+ let schema =
crate::arrow::build_target_arrow_schema(table.schema().fields())?;
+ for column in columns {
+ schema
+ .field_with_name(column)
+ .map_err(|error| invalid(error.to_string()))?;
+ }
+ // Schema order makes every projected batch interchangeable. An empty
+ // snapshot still carries its declared schema without fetching any
file.
+ let indices = schema
+ .fields()
+ .iter()
+ .enumerate()
+ .filter_map(|(index, field)|
columns.contains(field.name()).then_some(index))
+ .collect::<Vec<_>>();
+ let projected = std::sync::Arc::new(
+ schema
+ .project(&indices)
+ .map_err(|error| invalid(error.to_string()))?,
+ );
+ let names = projected
+ .fields()
+ .iter()
+ .map(|field| field.name().as_str())
+ .collect::<Vec<_>>();
+ let snapshot = super::time_travel::resolve_snapshot(table).await?;
+ // Pin the source independently, keeping branch/time-travel resolution
+ // and the caller's resolved schema. Include L0 as Python
plan_for_write.
+ let table = snapshot.as_ref().map_or_else(
+ || table.clone(),
+ |snapshot| {
+ table
+ .copy_with_pinned_snapshot(snapshot)
+ .copy_with_options(HashMap::from([("scan.mode".into(),
"default".into())]))
+ },
+ );
+ let mut builder = table.new_read_builder();
+ builder.with_projection(&names)?;
+ let plan = if snapshot.is_some() {
+ builder.new_scan().with_scan_all_files().plan().await?
+ } else {
+ // Run the normal reader's policy checks even with no snapshot, but
+ // do not observe files from a commit arriving after resolution.
+ super::source::Plan::new(Vec::new())
+ };
+ let mut stream = builder.new_read()?.to_arrow(plan.splits())?;
+ let mut batches = Vec::new();
+ while let Some(batch) = stream.try_next().await? {
+ if batch.num_rows() > 0 {
+ batches.push(batch);
+ }
+ }
+ if batches.is_empty() {
+ batches.push(RecordBatch::new_empty(projected));
+ }
+ Self::new(batches)
+ }
+
+ pub fn row_offset(&self, batch: usize) -> usize {
+ if batch == 0 {
+ 0
+ } else {
+ self.ends[batch - 1]
+ }
+ }
+
+ pub fn gather(&self, rows: &[usize]) -> crate::Result<RecordBatch> {
+ if rows.is_empty() {
+ return Ok(RecordBatch::new_empty(self.schema.clone()));
+ }
+ let mut used = Vec::new();
+ let mut positions = HashMap::new();
+ let mut indices = Vec::with_capacity(rows.len());
+ for &row in rows {
+ if row >= self.row_count {
+ return Err(invalid("MERGE source row is out of range"));
+ }
+ let batch = self.ends.partition_point(|end| *end <= row);
+ let position = *positions.entry(batch).or_insert_with(|| {
+ used.push(&self.batches[batch]);
+ used.len() - 1
+ });
+ indices.push((position, row - self.row_offset(batch)));
+ }
+ // Touch only contributing chunks, and preserve the requested order.
+ // Avoid concatenating entire source columns with 32-bit Arrow offsets.
+ interleave_record_batch(&used, &indices).map_err(|error|
invalid(error.to_string()))
+ }
+}
diff --git a/crates/paimon/src/table/table_merge_into.rs
b/crates/paimon/src/table/table_merge_into.rs
index e788867b..1eaae977 100644
--- a/crates/paimon/src/table/table_merge_into.rs
+++ b/crates/paimon/src/table/table_merge_into.rs
@@ -20,10 +20,11 @@
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
+use super::table_merge_input::MergeInput;
use arrow_array::{Array, BooleanArray, Int64Array, RecordBatch, UInt64Array};
use arrow_row::{RowConverter, SortField};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
-use arrow_select::{concat::concat_batches, take::take};
+use arrow_select::take::take;
use futures::{future::BoxFuture, TryStreamExt};
use super::data_evolution_writer::RowIdFileIndex;
@@ -68,9 +69,11 @@ pub struct WhenNotMatched {
pub assignments: Vec<(String, MergeAssignment)>,
}
-/// Already materialized input, or the target itself joined on its row ID.
+/// Materialized input, a Paimon source table, or the target joined on its row
ID.
pub enum MergeSource {
Batches(Vec<RecordBatch>),
+ /// Read a Paimon table at one independently pinned source snapshot.
+ Table(Arc<Table>),
SelfTable,
}
@@ -307,32 +310,49 @@ pub(super) async fn merge_into(
};
let index = RowIdFileIndex::from_splits(scan_table.clone(),
plan.splits())?;
let source = match source {
- MergeSource::Batches(batches) => {
- let first = batches
- .first()
- .ok_or_else(|| invalid("MERGE source needs an Arrow schema"))?;
- if batches.iter().any(|batch| batch.schema() != first.schema()) {
- return Err(invalid("MERGE source batches must share a
schema"));
- }
- let names = first
- .schema()
- .fields()
+ MergeSource::Batches(batches) => Some(MergeInput::new(batches)?),
+ MergeSource::Table(source) => {
+ let mut columns = on
.iter()
- .map(|field| field.name().clone())
+ .map(|(_, source)| source.clone())
.collect::<HashSet<_>>();
- if names.len() != first.num_columns() {
- return Err(invalid("MERGE source has duplicate column names"));
+ for (assignments, condition) in matched
+ .iter()
+ .map(|clause| (&clause.assignments, &clause.condition))
+ .chain(
+ not_matched
+ .iter()
+ .map(|clause| (&clause.assignments,
&clause.condition)),
+ )
+ {
+ for (_, assignment) in assignments {
+ if let MergeAssignment::SourceColumn(name) = assignment {
+ columns.insert(name.clone());
+ }
+ if matches!(
+ assignment,
+ MergeAssignment::Value(UpdateAssignment::Function(_))
+ ) {
+ columns.extend(
+ source
+ .schema()
+ .fields()
+ .iter()
+ .map(|field| field.name().to_string()),
+ );
+ }
+ }
+ if let Some(condition) = condition {
+ columns.extend(condition.source_columns.clone());
+ }
}
- Some(
- concat_batches(&first.schema(), &batches)
- .map_err(|error| invalid(error.to_string()))?,
- )
+ Some(MergeInput::read(&source, &columns).await?)
}
MergeSource::SelfTable => None,
};
let source_schema = source
.as_ref()
- .map_or_else(|| schema.clone(), RecordBatch::schema);
+ .map_or_else(|| schema.clone(), |source| source.schema.clone());
for (assignments, condition, insert) in matched
.iter()
.map(|clause| (&clause.assignments, &clause.condition, false))
@@ -390,14 +410,16 @@ pub(super) async fn merge_into(
}
let mut converter = None;
let mut source_keys = HashMap::<Vec<u8>, Vec<usize>>::new();
- let mut seen_source = vec![false; source.as_ref().map_or(0,
RecordBatch::num_rows)];
+ let mut seen_source = vec![false; source.as_ref().map_or(0, |source|
source.row_count)];
if let Some(source) = &source {
- let keys = on
+ let fields = on
.iter()
.map(|(target, source_name)| {
- let source_column = source
- .column_by_name(source_name)
- .ok_or_else(|| invalid(format!("Missing source ON key
{source_name}")))?;
+ let source_type = source
+ .schema
+ .field_with_name(source_name)
+ .map_err(|error| invalid(error.to_string()))?
+ .data_type();
let target_type = if target == ROW_ID {
&DataType::Int64
} else {
@@ -406,31 +428,33 @@ pub(super) async fn merge_into(
.map_err(|error| invalid(error.to_string()))?
.data_type()
};
- if source_column.data_type() != target_type
+ if source_type != target_type
||
!super::upsert_key_matcher::supported_key_type(target_type)
{
return Err(invalid(
"MERGE ON key types must match and support row
encoding",
));
}
- Ok(source_column.clone())
+ Ok(SortField::new(source_type.clone()))
})
- .collect::<crate::Result<Vec<_>>>()?;
- let key_converter = RowConverter::new(
- keys.iter()
- .map(|column| SortField::new(column.data_type().clone()))
- .collect(),
- )
- .map_err(|error| invalid(error.to_string()))?;
- let rows = key_converter
- .convert_columns(&keys)
- .map_err(|error| invalid(error.to_string()))?;
- for row in 0..source.num_rows() {
- if keys.iter().all(|column| !column.is_null(row)) {
- source_keys
- .entry(rows.row(row).as_ref().to_vec())
- .or_default()
- .push(row);
+ .collect::<crate::Result<_>>()?;
+ let key_converter =
+ RowConverter::new(fields).map_err(|error|
invalid(error.to_string()))?;
+ for (chunk, batch) in source.batches.iter().enumerate() {
+ let keys = on
+ .iter()
+ .map(|(_, name)| batch.column_by_name(name).unwrap().clone())
+ .collect::<Vec<_>>();
+ let rows = key_converter
+ .convert_columns(&keys)
+ .map_err(|error| invalid(error.to_string()))?;
+ for row in 0..batch.num_rows() {
+ if keys.iter().all(|column| !column.is_null(row)) {
+ source_keys
+ .entry(rows.row(row).as_ref().to_vec())
+ .or_default()
+ .push(source.row_offset(chunk) + row);
+ }
}
}
converter = Some(key_converter);
@@ -474,7 +498,7 @@ pub(super) async fn merge_into(
if !target_rows.is_empty() {
joined.push(aliases(
Some(&gather(&target, &target_rows)?),
- &gather(source, &source_rows)?,
+ &source.gather(&source_rows)?,
)?);
}
} else if target.num_rows() > 0 {
@@ -523,26 +547,36 @@ pub(super) async fn merge_into(
}
let mut inserts = Vec::new();
if let Some(source) = source {
- let unmatched = seen_source
- .iter()
- .enumerate()
- .filter_map(|(row, seen)| (!seen).then_some(row))
- .collect::<Vec<_>>();
- let mut pending = aliases(None, &gather(&source, &unmatched)?)?;
+ let mut pending = Vec::new();
+ for (chunk, batch) in source.batches.iter().enumerate() {
+ let offset = source.row_offset(chunk);
+ let unmatched = (0..batch.num_rows())
+ .filter(|row| !seen_source[offset + row])
+ .collect::<Vec<_>>();
+ if !unmatched.is_empty() {
+ pending.push(aliases(None, &gather(batch, &unmatched)?)?);
+ }
+ }
let columns = schema
.fields()
.iter()
.map(|field| field.name().clone())
.collect::<Vec<_>>();
for clause in not_matched {
- if pending.num_rows() == 0 {
- break;
+ let mut remaining = Vec::new();
+ let mut clause_rows = Vec::new();
+ for batch in pending {
+ let (selected, rest) = select_clause(batch,
&clause.condition).await?;
+ if rest.num_rows() > 0 {
+ remaining.push(rest);
+ }
+ if selected.num_rows() > 0 {
+ clause_rows.push(selected);
+ }
}
- let (selected, rest) = select_clause(pending,
&clause.condition).await?;
- pending = rest;
- if selected.num_rows() > 0 {
+ if !clause_rows.is_empty() {
for batch in assigned(
- std::slice::from_ref(&selected),
+ &clause_rows,
&clause.assignments,
&columns,
schema.clone(),
@@ -555,6 +589,7 @@ pub(super) async fn merge_into(
);
}
}
+ pending = remaining;
}
}
// Materialize/validate every action before staging. Once prepared,
messages
diff --git a/crates/paimon/src/table/table_merge_into_tests.rs
b/crates/paimon/src/table/table_merge_into_tests.rs
index 4b4ec6db..b20c4753 100644
--- a/crates/paimon/src/table/table_merge_into_tests.rs
+++ b/crates/paimon/src/table/table_merge_into_tests.rs
@@ -28,6 +28,10 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
fn table() -> Table {
+ table_at("merge")
+}
+
+fn table_at(name: &str) -> Table {
let schema = Schema::builder()
.column("id", DataType::Int(IntType::new()))
.column("value", DataType::Int(IntType::new()))
@@ -38,8 +42,8 @@ fn table() -> Table {
.unwrap();
Table::new(
FileIOBuilder::new("memory").build().unwrap(),
- Identifier::new("db", "t"),
- "memory:/merge".into(),
+ Identifier::new("db", name),
+ format!("memory:/{name}"),
TableSchema::new(0, &schema),
None,
)
@@ -475,3 +479,218 @@ async fn
core_merge_evaluates_functions_once_for_the_complete_clause() {
vec![(Some(1), Some(100)), (Some(2), Some(200))]
);
}
+
+#[tokio::test]
+async fn core_merge_reads_table_source_at_the_selected_snapshot() {
+ let target = table();
+ let source = table_at("source");
+ seed(&target, &batch(vec![Some(1)], vec![Some(10)])).await;
+ seed(&source, &batch(vec![Some(1)], vec![Some(11)])).await;
+ seed(&source, &batch(vec![Some(2)], vec![Some(22)])).await;
+ let historical =
source.copy_with_options(std::collections::HashMap::from([(
+ "scan.snapshot-id".into(),
+ "1".into(),
+ )]));
+ let messages = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Table(Arc::new(historical)),
+ vec![("id".into(), "id".into())],
+ vec![update(None)],
+ vec![WhenNotMatched {
+ condition: None,
+ assignments: vec![
+ ("id".into(), MergeAssignment::SourceColumn("id".into())),
+ (
+ "value".into(),
+ MergeAssignment::SourceColumn("value".into()),
+ ),
+ ],
+ }],
+ )
+ .await
+ .unwrap();
+ TableCommit::new(target.clone(), "merge".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(rows(&target).await, vec![(Some(1), Some(11))]);
+ assert_eq!(
+ rows(&source).await,
+ vec![(Some(1), Some(11)), (Some(2), Some(22))]
+ );
+}
+
+#[tokio::test]
+async fn core_merge_empty_table_source_keeps_no_match_assignments_lazy() {
+ let target = table();
+ seed(&target, &batch(vec![Some(1)], vec![Some(10)])).await;
+ let source = table_at("empty_source");
+ let called = Arc::new(AtomicUsize::new(0));
+ let evaluated = called.clone();
+ let value = UpdateAssignment::DeferredScalar(Arc::new(move || {
+ evaluated.fetch_add(1, Ordering::Relaxed);
+ Err(crate::Error::DataInvalid {
+ message: "Must remain lazy".into(),
+ source: None,
+ })
+ }));
+ let messages = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Table(Arc::new(source)),
+ vec![("id".into(), "id".into())],
+ vec![WhenMatched {
+ condition: None,
+ delete: false,
+ assignments: vec![("value".into(),
MergeAssignment::Value(value))],
+ }],
+ Vec::new(),
+ )
+ .await
+ .unwrap();
+ assert!(messages.is_empty());
+ assert_eq!(called.load(Ordering::Relaxed), 0);
+ assert_eq!(rows(&target).await, vec![(Some(1), Some(10))]);
+}
+
+#[tokio::test]
+async fn core_merge_empty_table_source_obeys_reader_authorization() {
+ let target = table();
+ seed(&target, &batch(vec![Some(1)], vec![Some(10)])).await;
+ let source =
table_at("empty_source").copy_with_options(std::collections::HashMap::from([(
+ "query-auth.enabled".into(),
+ "true".into(),
+ )]));
+ let error = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Table(Arc::new(source)),
+ vec![("id".into(), "id".into())],
+ vec![update(None)],
+ Vec::new(),
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("query-auth.enabled"));
+ assert_eq!(rows(&target).await, vec![(Some(1), Some(10))]);
+}
+
+#[tokio::test]
+async fn core_merge_checks_source_duplicates_across_batches() {
+ let target = table();
+ seed(&target, &batch(vec![Some(1)], vec![Some(10)])).await;
+ let source = vec![
+ batch(vec![Some(1)], vec![Some(11)]),
+ batch(vec![Some(1)], vec![Some(12)]),
+ ];
+ let error = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Batches(source),
+ vec![("id".into(), "id".into())],
+ vec![update(None)],
+ Vec::new(),
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("multiple source rows"));
+ assert_eq!(rows(&target).await, vec![(Some(1), Some(10))]);
+}
+
+#[tokio::test]
+async fn core_merge_does_not_reread_a_source_advanced_by_a_condition() {
+ let target = table();
+ let source = Arc::new(table_at("source_advanced"));
+ seed(&target, &batch(vec![Some(1)], vec![Some(10)])).await;
+ seed(&source, &batch(vec![Some(1)], vec![Some(11)])).await;
+ let advancing = source.clone();
+ let condition = MergeCondition {
+ target_columns: Vec::new(),
+ source_columns: Vec::new(),
+ evaluate: Arc::new(move |input| {
+ let source = advancing.clone();
+ Box::pin(async move {
+ seed(&source, &batch(vec![Some(2)], vec![Some(22)])).await;
+ Ok(BooleanArray::from(vec![true; input.num_rows()]))
+ })
+ }),
+ };
+ let messages = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Table(source.clone()),
+ vec![("id".into(), "id".into())],
+ vec![update(Some(condition))],
+ vec![WhenNotMatched {
+ condition: None,
+ assignments: vec![
+ ("id".into(), MergeAssignment::SourceColumn("id".into())),
+ (
+ "value".into(),
+ MergeAssignment::SourceColumn("value".into()),
+ ),
+ ],
+ }],
+ )
+ .await
+ .unwrap();
+ TableCommit::new(target.clone(), "merge".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(rows(&target).await, vec![(Some(1), Some(11))]);
+ assert_eq!(
+ rows(&source).await,
+ vec![(Some(1), Some(11)), (Some(2), Some(22))]
+ );
+}
+
+#[tokio::test]
+async fn core_merge_insert_arrays_cover_selected_rows_across_source_batches() {
+ let target = table();
+ let messages = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Batches(vec![
+ batch(vec![Some(1)], vec![Some(11)]),
+ batch(vec![Some(2)], vec![Some(22)]),
+ ]),
+ vec![("id".into(), "id".into())],
+ Vec::new(),
+ vec![WhenNotMatched {
+ condition: None,
+ assignments: vec![
+ ("id".into(), MergeAssignment::SourceColumn("id".into())),
+ (
+ "value".into(),
+
MergeAssignment::Value(UpdateAssignment::Array(vec![Arc::new(
+ Int32Array::from(vec![100, 200]),
+ )])),
+ ),
+ ],
+ }],
+ )
+ .await
+ .unwrap();
+ TableCommit::new(target.clone(), "merge".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(
+ rows(&target).await,
+ vec![(Some(1), Some(100)), (Some(2), Some(200))]
+ );
+}
