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 ca0f40d1 [python] Use typed MERGE clauses and named execution
factories (#1204)
ca0f40d1 is described below
commit ca0f40d1e521691c43bcdec03b9f37dbfc52beae
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 09:55:25 2026 +0800
[python] Use typed MERGE clauses and named execution factories (#1204)
---
bindings/python/project-description.md | 43 +++
.../python/python/pypaimon_rust/datafusion.pyi | 66 +++-
bindings/python/src/context.rs | 3 +
bindings/python/src/lib.rs | 1 +
bindings/python/src/merge_execution.rs | 382 +++++++++++++++------
bindings/python/src/merge_into.rs | 173 ++++++----
bindings/python/src/merge_types.rs | 228 ++++++++++++
bindings/python/src/update_assignment.rs | 129 +++----
bindings/python/src/write.rs | 216 +++++++++++-
bindings/python/tests/test_blob_uri_reader.py | 8 +-
bindings/python/tests/test_update.py | 30 +-
11 files changed, 1014 insertions(+), 265 deletions(-)
diff --git a/bindings/python/project-description.md
b/bindings/python/project-description.md
index 86c40860..9738ab17 100644
--- a/bindings/python/project-description.md
+++ b/bindings/python/project-description.md
@@ -85,6 +85,49 @@ updates, including incremental row-ID updaters, expose the
same chainable setter
`with_commit_user` is available on batch, stream, and postpone fixed bucket
write
builders, and the identity applies to both writers and committers from that
builder.
+### Typed MERGE Actions
+
+Batch and stream `merge_into` accept ordered `WhenMatched` and `WhenNotMatched`
+objects. Assignments use `MergeAssignment.source_column`, `target_column`,
+`literal`, or `function`. A function receives the selected clause's Arrow table
+and returns an Array or ChunkedArray with the same number of rows.
+
+Execution engines that already joined target and source data can reuse these
+types through the batch execution factories:
+
+```python
+import pyarrow as pa
+from pypaimon_rust.datafusion import MergeAssignment, MergeBatchExecution,
WhenMatched
+
+joined = pa.table({
+ "t._ROW_ID": [10, 11],
+ "t.value": pa.array([20, 30], type=pa.int32()),
+ "s.value": pa.array([40, 10], type=pa.int32()),
+})
+matched = [
+ WhenMatched.update(
+ {"value": MergeAssignment.source_column("value")},
+ condition='"s.value" > "t.value"',
+ ),
+ WhenMatched.delete(),
+]
+update_schema = pa.schema([
+ pa.field("_ROW_ID", pa.int64(), nullable=False),
+ pa.field("value", pa.int32()),
+])
+updates = MergeBatchExecution.for_updates(joined.schema, update_schema,
matched)
+deletes = MergeBatchExecution.for_deletes(joined.schema, matched)
+assert updates.execute(joined).to_pylist() == [{"_ROW_ID": 10, "value": 40}]
+assert deletes.execute(joined).to_pylist() == [{"_ROW_ID": 11}]
+```
+
+Both matched factories consume the complete ordered branch list. The delete
+factory skips UPDATE assignments and callbacks while preserving their
conditions.
+For unmatched source batches, use `for_inserts(source_schema, target_schema,
+when_not_matched)` with `WhenNotMatched` objects. The factories compile
conditions
+once and return an executor with `execute(Table | RecordBatch) -> Table`.
+The calling engine owns joins, authorization and table writes.
+
### Multimodal SQL Helpers
`SQLContext` also registers Python scalar UDFs for common BLOB media and vector
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index ae393712..976c3309 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -18,7 +18,7 @@
from datetime import datetime
from os import PathLike
from typing import (
- Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple,
+ Any, Callable, Dict, Iterable, List, Literal, Mapping, Optional, Sequence,
Tuple,
TypeAlias, Union,
)
@@ -79,6 +79,40 @@ class RowIdPlan:
...
def plan(self) -> Plan: ...
+class MergeAssignment:
+ """An immutable source reference, target reference, literal or callback."""
+ @staticmethod
+ def source_column(name: str) -> "MergeAssignment": ...
+ @staticmethod
+ def target_column(name: str) -> "MergeAssignment": ...
+ @staticmethod
+ def literal(value: Any) -> "MergeAssignment":
+ """Defer scalar conversion; Arrow arrays provide values for selected
rows."""
+ ...
+ @staticmethod
+ def function(
+ function: Callable[[pyarrow.Table], Union[pyarrow.Array,
pyarrow.ChunkedArray]],
+ ) -> "MergeAssignment": ...
+
+class WhenMatched:
+ """An immutable, ordered UPDATE or DELETE clause with an optional SQL
condition.
+
+ Conditions reference quoted Arrow names, such as ``"s.value" > "t.value"``.
+ The first matching clause consumes a row; NULL conditions fall through.
+ """
+ @staticmethod
+ def update(
+ assignments: Mapping[str, MergeAssignment], condition: Optional[str] =
None,
+ ) -> "WhenMatched": ...
+ @staticmethod
+ def delete(condition: Optional[str] = None) -> "WhenMatched": ...
+
+class WhenNotMatched:
+ """An immutable, ordered INSERT clause."""
+ def __init__(
+ self, assignments: Mapping[str, MergeAssignment], condition:
Optional[str] = None,
+ ) -> None: ...
+
class MergeBatchExecution:
"""Evaluate ordered MERGE actions over already joined Arrow batches.
@@ -89,10 +123,22 @@ class MergeBatchExecution:
output contains ``_ROW_ID`` and the requested target columns; delete output
contains only ``_ROW_ID``. Insert output uses the target schema.
"""
- def __init__(
- self, input_schema: pyarrow.Schema, output_schema: pyarrow.Schema,
- clauses: Sequence[Dict[str, Any]], mode: Literal["update", "delete",
"insert"],
- ) -> None: ...
+ @staticmethod
+ def for_updates(
+ input_schema: pyarrow.Schema, update_schema: pyarrow.Schema,
+ when_matched: Sequence[WhenMatched],
+ ) -> "MergeBatchExecution": ...
+ @staticmethod
+ def for_deletes(
+ input_schema: pyarrow.Schema, when_matched: Sequence[WhenMatched],
+ ) -> "MergeBatchExecution":
+ """Keep all matched conditions; skip UPDATE assignment conversion and
evaluation."""
+ ...
+ @staticmethod
+ def for_inserts(
+ input_schema: pyarrow.Schema, target_schema: pyarrow.Schema,
+ when_not_matched: Sequence[WhenNotMatched],
+ ) -> "MergeBatchExecution": ...
def execute(self, batch: Union[pyarrow.Table, pyarrow.RecordBatch]) ->
pyarrow.Table:
"""Return selected actions; unselected assignments are never
evaluated."""
...
@@ -542,15 +588,14 @@ class BatchTableUpdate:
) -> "BatchTableUpdate": ...
def merge_into(
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]]
+ when_matched: Sequence[WhenMatched], when_not_matched:
Sequence[WhenNotMatched]
) -> List[CommitMessage]:
"""Prepare MERGE against one target snapshot. None selects self-merge
on _ROW_ID.
A Table source reads its selected snapshot or numeric timestamp window
in core, independently of the target. Arrow input preserves batches
and source row order.
- Clauses carry assignments as (column, source/target/literal/callable,
value)
- tuples, a delete flag, and optional SQL condition metadata. Callables
+ Typed clauses carry MergeAssignment values and optional SQL
conditions. Callables
consume the complete clause's qualified t./s. Arrow columns and return
an Array or ChunkedArray with the same row count; empty clauses are
lazy.
Commit returned messages through the matching builder's committer.
@@ -616,7 +661,7 @@ class StreamTableUpdate:
) -> "StreamTableUpdate": ...
def merge_into(
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]],
+ when_matched: Sequence[WhenMatched], when_not_matched:
Sequence[WhenNotMatched],
commit_identifier: int
) -> List[CommitMessage]:
"""Prepare MERGE against one target snapshot. None selects self-merge
on _ROW_ID.
@@ -624,8 +669,7 @@ class StreamTableUpdate:
A Table source reads its selected snapshot or numeric timestamp window
in core, independently of the target. Arrow input preserves batches
and source row order.
- Clauses carry assignments as (column, source/target/literal/callable,
value)
- tuples, a delete flag, and optional SQL condition metadata. Callables
+ Typed clauses carry MergeAssignment values and optional SQL
conditions. Callables
consume the complete clause's qualified t./s. Arrow columns and return
an Array or ChunkedArray with the same row count; empty clauses are
lazy.
Commit returned messages through the matching builder's committer.
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 162c0444..112a5014 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -1039,6 +1039,9 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<crate::write::PyBatchTableUpdate>()?;
this.add_class::<crate::write::PyTableUpdateByRowId>()?;
this.add_class::<crate::merge_execution::PyMergeBatchExecution>()?;
+ this.add_class::<crate::merge_types::PyMergeAssignment>()?;
+ this.add_class::<crate::merge_types::PyWhenMatched>()?;
+ this.add_class::<crate::merge_types::PyWhenNotMatched>()?;
this.add_class::<crate::shard_update::PyShardTableUpdator>()?;
this.add_class::<crate::write::PyStreamTableUpdate>()?;
this.add_class::<crate::write::PyStreamTableWrite>()?;
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index e45485bf..69f081bb 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -31,6 +31,7 @@ mod hybrid_search;
mod index_build;
mod merge_execution;
mod merge_into;
+mod merge_types;
mod oss_cpp;
mod predicate;
mod read;
diff --git a/bindings/python/src/merge_execution.rs
b/bindings/python/src/merge_execution.rs
index 7310bfdc..7d4f17db 100644
--- a/bindings/python/src/merge_execution.rs
+++ b/bindings/python/src/merge_execution.rs
@@ -21,7 +21,7 @@ use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use arrow::array::Int64Array;
-use arrow::datatypes::{DataType, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use arrow::record_batch::RecordBatch;
use paimon::table::{
@@ -30,7 +30,7 @@ use paimon::table::{
use paimon_datafusion::runtime::runtime;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList};
+use pyo3::types::PyList;
use crate::error::to_py_err;
@@ -79,28 +79,14 @@ fn row_ids(batch: RecordBatch) -> PyResult<RecordBatch> {
RecordBatch::try_new(schema, columns).map_err(|error|
PyValueError::new_err(error.to_string()))
}
-#[pymethods]
impl PyMergeBatchExecution {
- #[new]
- fn new(
+ fn compile(
py: Python<'_>,
- input_schema: &Bound<'_, PyAny>,
- output_schema: &Bound<'_, PyAny>,
+ input_schema: SchemaRef,
+ output_schema: SchemaRef,
clauses: &Bound<'_, PyAny>,
- mode: &str,
+ output: MergeActionOutput,
) -> PyResult<Self> {
- let output = match mode {
- "update" => MergeActionOutput::Updates,
- "delete" => MergeActionOutput::Deletes,
- "insert" => MergeActionOutput::Inserts,
- _ => {
- return Err(PyValueError::new_err(
- "MERGE execution mode must be update, delete or insert",
- ));
- }
- };
- let input_schema = Arc::new(Schema::from_pyarrow_bound(input_schema)?);
- let output_schema =
Arc::new(Schema::from_pyarrow_bound(output_schema)?);
if output != MergeActionOutput::Inserts {
let row_id = output_schema.field_with_name("_ROW_ID").map_err(|_| {
PyValueError::new_err("Matched MERGE output needs an Int64
_ROW_ID field")
@@ -110,60 +96,46 @@ impl PyMergeBatchExecution {
"Matched MERGE output needs an Int64 _ROW_ID field",
));
}
- if output == MergeActionOutput::Deletes &&
output_schema.fields().len() != 1 {
- return Err(PyValueError::new_err(
- "MERGE DELETE output must contain only _ROW_ID",
- ));
- }
}
let condition_schema = if output == MergeActionOutput::Inserts {
input_schema.clone()
} else {
matched_schema(input_schema.clone())?
};
- let target_schema = Arc::new(Schema::new(
+ let target_schema = Arc::new(Schema::new_with_metadata(
output_schema
.fields()
.iter()
.filter(|field| output == MergeActionOutput::Inserts ||
field.name() != "_ROW_ID")
.cloned()
.collect::<Vec<_>>(),
+ output_schema.metadata().clone(),
));
let callback_error = Arc::new(Mutex::new(None));
- let mut matched = Vec::new();
- let mut not_matched = Vec::new();
- for clause in clauses.try_iter()? {
- let clause = clause?;
- let clause = clause.cast::<PyDict>()?;
- let condition = crate::merge_into::condition(py, clause,
condition_schema.clone())?;
- let assignments = if output == MergeActionOutput::Deletes {
- // A DELETE route still consumes preceding UPDATE conditions,
- // but must not inspect or evaluate those assignments.
- Vec::new()
- } else {
- crate::merge_into::assignments(
+ let (mut matched, not_matched) = if output ==
MergeActionOutput::Inserts {
+ (
+ Vec::new(),
+ crate::merge_into::not_matched(
py,
- clause,
+ clauses,
+ condition_schema,
target_schema.clone(),
callback_error.clone(),
- )?
- };
- if output == MergeActionOutput::Inserts {
- not_matched.push(WhenNotMatched {
- condition,
- assignments,
- });
- } else {
- matched.push(WhenMatched {
- condition,
- delete: clause
- .get_item("delete")?
- .ok_or_else(|| PyValueError::new_err("MERGE matched
clause needs action"))?
- .extract()?,
- assignments,
- });
- }
- }
+ )?,
+ )
+ } else {
+ (
+ crate::merge_into::matched(
+ py,
+ clauses,
+ condition_schema,
+ target_schema.clone(),
+ callback_error.clone(),
+ output != MergeActionOutput::Deletes,
+ )?,
+ Vec::new(),
+ )
+ };
if output == MergeActionOutput::Updates {
// The explicit output schema may request fields untouched by every
// clause. Include them in core's shared update-column projection.
@@ -206,6 +178,62 @@ impl PyMergeBatchExecution {
execution: Mutex::new(()),
})
}
+}
+
+#[pymethods]
+impl PyMergeBatchExecution {
+ #[staticmethod]
+ fn for_updates(
+ py: Python<'_>,
+ input_schema: &Bound<'_, PyAny>,
+ update_schema: &Bound<'_, PyAny>,
+ when_matched: &Bound<'_, PyAny>,
+ ) -> PyResult<Self> {
+ Self::compile(
+ py,
+ Arc::new(Schema::from_pyarrow_bound(input_schema)?),
+ Arc::new(Schema::from_pyarrow_bound(update_schema)?),
+ when_matched,
+ MergeActionOutput::Updates,
+ )
+ }
+
+ #[staticmethod]
+ fn for_deletes(
+ py: Python<'_>,
+ input_schema: &Bound<'_, PyAny>,
+ when_matched: &Bound<'_, PyAny>,
+ ) -> PyResult<Self> {
+ let input_schema = Arc::new(Schema::from_pyarrow_bound(input_schema)?);
+ let output_schema = Arc::new(Schema::new(vec![Field::new(
+ "_ROW_ID",
+ DataType::Int64,
+ false,
+ )]));
+ Self::compile(
+ py,
+ input_schema,
+ output_schema,
+ when_matched,
+ MergeActionOutput::Deletes,
+ )
+ }
+
+ #[staticmethod]
+ fn for_inserts(
+ py: Python<'_>,
+ input_schema: &Bound<'_, PyAny>,
+ target_schema: &Bound<'_, PyAny>,
+ when_not_matched: &Bound<'_, PyAny>,
+ ) -> PyResult<Self> {
+ Self::compile(
+ py,
+ Arc::new(Schema::from_pyarrow_bound(input_schema)?),
+ Arc::new(Schema::from_pyarrow_bound(target_schema)?),
+ when_not_matched,
+ MergeActionOutput::Inserts,
+ )
+ }
fn execute(&self, py: Python<'_>, batch: &Bound<'_, PyAny>) ->
PyResult<Py<PyAny>> {
let schema = Schema::from_pyarrow_bound(&batch.getattr("schema")?)?;
@@ -302,6 +330,7 @@ impl PyMergeBatchExecution {
mod tests {
use super::*;
use pyo3::ffi::c_str;
+ use pyo3::types::PyDict;
fn locals(py: Python<'_>) -> Bound<'_, PyDict> {
let locals = PyDict::new(py);
@@ -312,6 +341,24 @@ mod tests {
)
.unwrap();
locals
+ .set_item(
+ "MergeAssignment",
+ py.get_type::<crate::merge_types::PyMergeAssignment>(),
+ )
+ .unwrap();
+ locals
+ .set_item(
+ "WhenMatched",
+ py.get_type::<crate::merge_types::PyWhenMatched>(),
+ )
+ .unwrap();
+ locals
+ .set_item(
+ "WhenNotMatched",
+ py.get_type::<crate::merge_types::PyWhenNotMatched>(),
+ )
+ .unwrap();
+ locals
}
#[test]
@@ -346,10 +393,11 @@ def update(rows):
assert rows['s.value'].to_pylist() == [15, 20]
return pa.chunked_array([pa.array([30]), pa.array([40])])
clauses = [
- dict(condition={'sql': '"s.value" >= 90'}, delete=True, assignments=[]),
- dict(condition={'sql': '"s.value" > 10'}, delete=False,
- assignments=[('value', 'callable', update), ('keep', 'target',
'keep')]),
- dict(condition=None, delete=False, assignments=[('value', 'literal', 50)]),
+ WhenMatched.delete(condition='"s.value" >= 90'),
+ WhenMatched.update({'value': MergeAssignment.function(update),
+ 'keep': MergeAssignment.target_column('keep')},
+ condition='"s.value" > 10'),
+ WhenMatched.update({'value': MergeAssignment.literal(50)}),
]
output_schema = pa.schema([
pa.field('_ROW_ID', pa.int64(), nullable=False),
@@ -357,7 +405,7 @@ output_schema = pa.schema([
pa.field('unassigned', pa.int32()),
pa.field('value', pa.int32()),
], metadata={'output': 'updates'})
-execution = MergeBatchExecution(input_schema, output_schema, clauses, 'update')
+execution = MergeBatchExecution.for_updates(input_schema, output_schema,
clauses)
result = execution.execute(joined)
assert result.schema.equals(output_schema, check_metadata=True)
assert result.to_pydict() == {
@@ -367,7 +415,7 @@ assert result.to_pydict() == {
}
assert calls == [[1, 4]]
delete_schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)])
-deletes = MergeBatchExecution(input_schema, delete_schema, clauses,
'delete').execute(joined)
+deletes = MergeBatchExecution.for_deletes(input_schema,
clauses).execute(joined)
assert deletes.schema == delete_schema
assert deletes.to_pydict() == {'_ROW_ID': [3]}
assert calls == [[1, 4]]
@@ -394,26 +442,29 @@ source = pa.table({'s.value': pa.array([None, 20, 5],
pa.int32()),
's.id': pa.array([3, 1, 2], pa.int32())})
source = pa.Table.from_batches(source.to_batches(max_chunksize=1))
target = pa.schema([pa.field('label', pa.string()), pa.field('id', pa.int32()),
- pa.field('value', pa.int32()), pa.field('missing',
pa.bool_())])
+ pa.field('value', pa.int32()), pa.field('missing',
pa.bool_())],
+ metadata={'output': 'inserts'})
calls = []
def insert(rows):
calls.append(rows['s.id'].to_pylist())
return pa.array([21], pa.int32())
clauses = [
- dict(condition={'sql': '"s.value" > 10'},
- assignments=[('id', 'source', 'id'), ('value', 'callable', insert),
- ('label', 'literal', 'first')]),
- dict(condition=None,
- assignments=[('id', 'source', 'id'), ('value', 'literal', 7),
- ('label', 'literal', 'last'), ('missing', 'literal',
True)]),
+ WhenNotMatched({'id': MergeAssignment.source_column('id'),
+ 'value': MergeAssignment.function(insert),
+ 'label': MergeAssignment.literal('first')},
+ condition='"s.value" > 10'),
+ WhenNotMatched({'id': MergeAssignment.source_column('id'),
+ 'value': MergeAssignment.literal(7),
+ 'label': MergeAssignment.literal('last'),
+ 'missing': MergeAssignment.literal(True)}),
]
-execution = MergeBatchExecution(source.schema, target, clauses, 'insert')
+execution = MergeBatchExecution.for_inserts(source.schema, target, clauses)
result = execution.execute(source)
-assert result.schema == target
+assert result.schema.equals(target, check_metadata=True)
assert result.to_pydict() == {'label': ['first', 'last', 'last'], 'id': [1, 3,
2],
'value': [21, 7, 7], 'missing': [None, True, True]}
assert calls == [[1]]
-assert execution.execute(source.slice(0, 0)).schema == target
+assert execution.execute(source.slice(0, 0)).schema.equals(target,
check_metadata=True)
assert calls == [[1]]
"#
),
@@ -443,9 +494,8 @@ def assignment(rows):
if len(attempts) == 1:
raise failure
return pa.array([9], pa.int32())
-clauses = [dict(condition=None, delete=False,
- assignments=[('value', 'callable', assignment)])]
-execution = MergeBatchExecution(joined.schema, output, clauses, 'update')
+clauses = [WhenMatched.update({'value': MergeAssignment.function(assignment)})]
+execution = MergeBatchExecution.for_updates(joined.schema, output, clauses)
try:
execution.execute(joined)
except LookupError as error:
@@ -484,31 +534,171 @@ def rejected(function, text):
assert text in str(error), str(error)
else:
raise AssertionError('invalid MERGE operation succeeded')
-clause = dict(condition={'sql': '"s.absent" > 0'}, delete=False,
- assignments=[('value', 'literal', 1)])
-rejected(lambda: MergeBatchExecution(joined.schema, output, [clause],
'update'), 'absent')
-clause = dict(condition=None, delete=False,
- assignments=[('value', 'callable', lambda rows: pa.array([1],
pa.int32()))])
-execution = MergeBatchExecution(joined.schema, output, [clause], 'update')
+clause = WhenMatched.update({'value': MergeAssignment.literal(1)},
+ condition='"s.absent" > 0')
+rejected(lambda: MergeBatchExecution.for_updates(joined.schema, output,
[clause]), 'absent')
+clause = WhenMatched.update({'value': MergeAssignment.function(
+ lambda rows: pa.array([1], pa.int32()))})
+execution = MergeBatchExecution.for_updates(joined.schema, output, [clause])
rejected(lambda: execution.execute(joined), 'length')
rejected(lambda: execution.execute(joined.rename_columns(['t._ROW_ID',
's.other'])), 'schema')
-rejected(lambda: MergeBatchExecution(joined.schema, output, [clause],
'unknown'), 'mode')
-rejected(lambda: MergeBatchExecution(joined.schema, pa.schema([('value',
pa.int32())]),
- [clause], 'update'), '_ROW_ID')
-rejected(lambda: MergeBatchExecution(joined.schema, pa.schema([('_ROW_ID',
pa.int32())]),
- [clause], 'delete'), 'Int64')
-rejected(lambda: MergeBatchExecution(joined.schema, output, [clause],
'delete'), 'only')
-rejected(lambda: MergeBatchExecution(joined.schema, output, [clause, clause],
'update'), 'last')
+rejected(lambda: MergeBatchExecution.for_updates(joined.schema,
+ pa.schema([('value', pa.int32())]), [clause]), '_ROW_ID')
+rejected(lambda: MergeBatchExecution.for_updates(joined.schema,
+ pa.schema([('_ROW_ID', pa.int32()), ('value', pa.int32())]), [clause]),
'Int64')
+rejected(lambda: MergeBatchExecution.for_updates(joined.schema, output,
[clause, clause]), 'last')
dual_ids = joined.append_column('_ROW_ID', pa.array([99, 98], pa.int64()))
-rejected(lambda: MergeBatchExecution(dual_ids.schema, output, [clause],
'update'), 'ambiguous')
-bad_type = dict(condition=None, delete=False,
- assignments=[('value', 'callable', lambda rows: [1, 2])])
-rejected(lambda: MergeBatchExecution(joined.schema, output, [bad_type],
'update').execute(joined),
+rejected(lambda: MergeBatchExecution.for_updates(dual_ids.schema, output,
[clause]), 'ambiguous')
+rejected(lambda: MergeBatchExecution.for_deletes(dual_ids.schema,
[WhenMatched.delete()]), 'ambiguous')
+rejected(lambda: MergeBatchExecution.for_deletes(pa.schema([('s.value',
pa.int32())]),
+ [WhenMatched.delete()]), '_ROW_ID')
+wrong_ids = pa.table({'t._ROW_ID': pa.array([1], pa.int32())})
+rejected(lambda: MergeBatchExecution.for_deletes(wrong_ids.schema,
+ [WhenMatched.delete()]).execute(wrong_ids), 'Int64')
+null_ids = pa.table({'_ROW_ID': pa.array([None], pa.int64())})
+rejected(lambda: MergeBatchExecution.for_deletes(null_ids.schema,
+ [WhenMatched.delete()]).execute(null_ids), 'NULL')
+bad_type = WhenMatched.update({'value': MergeAssignment.function(lambda rows:
[1, 2])})
+rejected(lambda: MergeBatchExecution.for_updates(joined.schema, output,
[bad_type]).execute(joined),
'pyarrow.Array')
# A no-match clause must not convert an unused literal or invoke a callback.
-clause = dict(condition={'sql': '"s.value" < 0'}, delete=False,
- assignments=[('value', 'literal', object())])
-assert MergeBatchExecution(joined.schema, output, [clause],
'update').execute(joined).num_rows == 0
+clause = WhenMatched.update({'value': MergeAssignment.literal(object())},
+ condition='"s.value" < 0')
+assert MergeBatchExecution.for_updates(joined.schema, output,
[clause]).execute(joined).num_rows == 0
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn named_factories_reject_dictionary_protocol_and_wrong_clause_families() {
+ Python::attach(|py| {
+ let locals = locals(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+input_schema = pa.schema([('t._ROW_ID', pa.int64()), ('s.value', pa.int32())])
+update_schema = pa.schema([('_ROW_ID', pa.int64()), ('value', pa.int32())])
+target_schema = pa.schema([('value', pa.int32())])
+matched = WhenMatched.update({'value': MergeAssignment.source_column('value')})
+not_matched = WhenNotMatched({'value': MergeAssignment.source_column('value')})
+legacy = dict(condition=None, delete=False,
+ assignments=[('value', 'source', 'value')])
+def rejected(function):
+ try:
+ function()
+ except TypeError:
+ pass
+ else:
+ raise AssertionError('untyped MERGE protocol or wrong clause family
was accepted')
+for clause in [legacy, not_matched]:
+ rejected(lambda: MergeBatchExecution.for_updates(input_schema,
update_schema, [clause]))
+ rejected(lambda: MergeBatchExecution.for_deletes(input_schema, [clause]))
+for clause in [legacy, matched]:
+ rejected(lambda: MergeBatchExecution.for_inserts(input_schema,
target_schema, [clause]))
+rejected(lambda: MergeBatchExecution(input_schema, update_schema, [legacy],
'update'))
+rejected(lambda: MergeBatchExecution())
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn deletes_keep_update_conditions_without_converting_update_assignments() {
+ Python::attach(|py| {
+ let locals = locals(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+schema = pa.schema([
+ pa.field('_ROW_ID', pa.int64(), nullable=True, metadata={'input':
'row-id'}),
+ pa.field('s.value', pa.int32()),
+], metadata={'input': 'joined'})
+joined = pa.record_batch([pa.array([1, 2, 3], pa.int64()),
+ pa.array([10, 20, 30], pa.int32())], schema=schema)
+calls = []
+def unused(rows):
+ calls.append(rows.num_rows)
+ raise AssertionError('DELETE evaluated an UPDATE callback')
+clauses = [
+ WhenMatched.update({'absent': MergeAssignment.literal(object()),
+ 'also_absent':
MergeAssignment.source_column('missing')},
+ condition='"_ROW_ID" = 1'),
+ WhenMatched.update({'unavailable': MergeAssignment.function(unused)},
+ condition='"s.value" > 20'),
+ WhenMatched.delete(),
+]
+execution = MergeBatchExecution.for_deletes(input_schema=schema,
when_matched=clauses)
+result = execution.execute(joined)
+expected = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)])
+assert result.schema.equals(expected, check_metadata=True)
+assert result.to_pydict() == {'_ROW_ID': [2]}
+assert execution.execute(schema.empty_table()).schema.equals(expected,
check_metadata=True)
+assert calls == []
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn cached_execution_keeps_concurrent_callback_errors_with_their_callers() {
+ Python::attach(|py| {
+ let locals = locals(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+import threading
+from concurrent.futures import ThreadPoolExecutor
+schema = pa.schema([('t._ROW_ID', pa.int64()), ('s.value', pa.int32())])
+output = pa.schema([('_ROW_ID', pa.int64()), ('value', pa.int32())])
+failures = {i: LookupError('assignment failed for ' + str(i)) for i in
range(4)}
+callback_entered = threading.Event()
+release_callback = threading.Event()
+def assignment(rows):
+ row_id = rows['t._ROW_ID'][0].as_py()
+ callback_entered.set()
+ assert release_callback.wait(timeout=10)
+ if row_id in failures:
+ raise failures[row_id]
+ return pa.array([9], pa.int32())
+execution = MergeBatchExecution.for_updates(schema, output,
+ [WhenMatched.update({'value': MergeAssignment.function(assignment)})])
+barrier = threading.Barrier(4)
+def execute(row_id):
+ batch = pa.record_batch([pa.array([row_id], pa.int64()), pa.array([1],
pa.int32())],
+ schema=schema)
+ barrier.wait()
+ try:
+ execution.execute(batch)
+ except LookupError as error:
+ return error
+ raise AssertionError('callback exception was lost')
+with ThreadPoolExecutor(max_workers=4) as threads:
+ futures = [threads.submit(execute, row_id) for row_id in range(4)]
+ try:
+ assert callback_entered.wait(timeout=10)
+ finally:
+ release_callback.set()
+ results = [future.result(timeout=10) for future in futures]
+assert all(error is failures[i] for i, error in enumerate(results))
+failures.clear()
+batch = pa.record_batch([pa.array([7], pa.int64()), pa.array([1],
pa.int32())], schema=schema)
+assert execution.execute(batch).to_pydict() == {'_ROW_ID': [7], 'value': [9]}
"#
),
Some(&locals),
diff --git a/bindings/python/src/merge_into.rs
b/bindings/python/src/merge_into.rs
index 3d6611ae..150bf36a 100644
--- a/bindings/python/src/merge_into.rs
+++ b/bindings/python/src/merge_into.rs
@@ -15,79 +15,126 @@
// specific language governing permissions and limitations
// under the License.
-//! Language transport and DataFusion expression compilation for core MERGE.
+//! Python MERGE descriptions and DataFusion expression compilation for core
MERGE.
use std::sync::{Arc, Mutex};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use paimon::table::{MergeAssignment, MergeCondition, WhenMatched,
WhenNotMatched};
use paimon_datafusion::runtime::runtime;
-use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
-use pyo3::types::PyDict;
use crate::error::to_py_err;
+use crate::merge_types::{AssignmentValue, Assignments, PyWhenMatched,
PyWhenNotMatched};
type CallbackError = Arc<Mutex<Option<PyErr>>>;
pub(crate) fn condition(
py: Python<'_>,
- clause: &Bound<'_, PyDict>,
+ sql: Option<&str>,
schema: SchemaRef,
) -> PyResult<Option<MergeCondition>> {
- let Some(value) = clause
- .get_item("condition")?
- .filter(|value| !value.is_none())
- else {
- return Ok(None);
- };
- let value = value.cast::<PyDict>()?;
- let sql: String = value
- .get_item("sql")?
- .ok_or_else(|| PyValueError::new_err("MERGE condition needs SQL"))?
- .extract()?;
- py.detach(||
runtime().block_on(paimon_datafusion::compile_merge_condition(sql, schema)))
- .map(Some)
+ sql.map(|sql| {
+ py.detach(|| {
+ runtime().block_on(paimon_datafusion::compile_merge_condition(
+ sql.to_owned(),
+ schema,
+ ))
+ })
.map_err(to_py_err)
+ })
+ .transpose()
}
pub(crate) fn assignments(
py: Python<'_>,
- clause: &Bound<'_, PyDict>,
+ values: &Assignments,
schema: SchemaRef,
errors: CallbackError,
) -> PyResult<Vec<(String, MergeAssignment)>> {
- let values = clause
- .get_item("assignments")?
- .ok_or_else(|| PyValueError::new_err("MERGE clause needs
assignments"))?;
values
- .try_iter()?
- .map(|item| {
- let (name, kind, value): (String, String, Py<PyAny>) =
item?.extract()?;
- let assignment = match kind.as_str() {
- "source" =>
MergeAssignment::SourceColumn(value.bind(py).extract()?),
- "target" =>
MergeAssignment::TargetColumn(value.bind(py).extract()?),
- "literal" | "callable" => {
- if kind == "callable" && !value.bind(py).is_callable() {
- return Err(PyValueError::new_err(
- "MERGE callable assignment must be callable",
- ));
- }
- let values = PyDict::new(py);
- values.set_item(&name, value)?;
- let assignment =
crate::update_assignment::from_python_schema(
+ .iter()
+ .map(|(name, value)| {
+ let value = value.borrow(py);
+ let assignment = match &value.value {
+ AssignmentValue::SourceColumn(column) => {
+ MergeAssignment::SourceColumn(column.clone())
+ }
+ AssignmentValue::TargetColumn(column) => {
+ MergeAssignment::TargetColumn(column.clone())
+ }
+ AssignmentValue::Literal(value) => {
+ MergeAssignment::Value(crate::update_assignment::literal(
py,
- &values,
+ name.clone(),
+ value.clone_ref(py),
schema.clone(),
errors.clone(),
- )?
- .remove(0)
- .1;
- MergeAssignment::Value(assignment)
+ )?)
}
- _ => return Err(PyValueError::new_err("Unknown MERGE
assignment kind")),
+ AssignmentValue::Function(value) => MergeAssignment::Value(
+ crate::update_assignment::function(value.clone_ref(py),
errors.clone()),
+ ),
};
- Ok((name, assignment))
+ Ok((name.clone(), assignment))
+ })
+ .collect()
+}
+
+pub(crate) fn matched(
+ py: Python<'_>,
+ clauses: &Bound<'_, PyAny>,
+ condition_schema: SchemaRef,
+ target_schema: SchemaRef,
+ errors: CallbackError,
+ include_assignments: bool,
+) -> PyResult<Vec<WhenMatched>> {
+ clauses
+ .try_iter()?
+ .map(|item| {
+ let item = item?;
+ let clause = item.extract::<PyRef<PyWhenMatched>>()?;
+ Ok(WhenMatched {
+ condition: condition(py, clause.condition.as_deref(),
condition_schema.clone())?,
+ delete: clause.delete,
+ assignments: if include_assignments {
+ assignments(
+ py,
+ &clause.assignments,
+ target_schema.clone(),
+ errors.clone(),
+ )?
+ } else {
+ // DELETE output consumes preceding UPDATE conditions
without
+ // inspecting values that this route will never
materialize.
+ Vec::new()
+ },
+ })
+ })
+ .collect()
+}
+
+pub(crate) fn not_matched(
+ py: Python<'_>,
+ clauses: &Bound<'_, PyAny>,
+ condition_schema: SchemaRef,
+ target_schema: SchemaRef,
+ errors: CallbackError,
+) -> PyResult<Vec<WhenNotMatched>> {
+ clauses
+ .try_iter()?
+ .map(|item| {
+ let item = item?;
+ let clause = item.extract::<PyRef<PyWhenNotMatched>>()?;
+ Ok(WhenNotMatched {
+ condition: condition(py, clause.condition.as_deref(),
condition_schema.clone())?,
+ assignments: assignments(
+ py,
+ &clause.assignments,
+ target_schema.clone(),
+ errors.clone(),
+ )?,
+ })
})
.collect()
}
@@ -96,8 +143,8 @@ pub(crate) fn from_python(
py: Python<'_>,
table: &paimon::table::Table,
source_schema: SchemaRef,
- matched: &Bound<'_, PyAny>,
- not_matched: &Bound<'_, PyAny>,
+ when_matched: &Bound<'_, PyAny>,
+ when_not_matched: &Bound<'_, PyAny>,
errors: CallbackError,
) -> PyResult<(Vec<WhenMatched>, Vec<WhenNotMatched>)> {
let target_schema =
@@ -115,22 +162,14 @@ pub(crate) fn from_python(
.clone()
.with_name(format!("s.{}", field.name()))
}));
- let schema = Arc::new(Schema::new(fields));
- let matched = matched
- .try_iter()?
- .map(|item| {
- let item = item?;
- let clause = item.cast::<PyDict>()?;
- Ok(WhenMatched {
- condition: condition(py, clause, schema.clone())?,
- delete: clause
- .get_item("delete")?
- .ok_or_else(|| PyValueError::new_err("MERGE matched clause
needs action"))?
- .extract()?,
- assignments: assignments(py, clause, target_schema.clone(),
errors.clone())?,
- })
- })
- .collect::<PyResult<Vec<_>>>()?;
+ let matched = matched(
+ py,
+ when_matched,
+ Arc::new(Schema::new(fields)),
+ target_schema.clone(),
+ errors.clone(),
+ true,
+ )?;
let schema = Arc::new(Schema::new(
source_schema
.fields()
@@ -143,16 +182,6 @@ pub(crate) fn from_python(
})
.collect::<Vec<_>>(),
));
- let not_matched = not_matched
- .try_iter()?
- .map(|item| {
- let item = item?;
- let clause = item.cast::<PyDict>()?;
- Ok(WhenNotMatched {
- condition: condition(py, clause, schema.clone())?,
- assignments: assignments(py, clause, target_schema.clone(),
errors.clone())?,
- })
- })
- .collect::<PyResult<Vec<_>>>()?;
+ let not_matched = not_matched(py, when_not_matched, schema, target_schema,
errors)?;
Ok((matched, not_matched))
}
diff --git a/bindings/python/src/merge_types.rs
b/bindings/python/src/merge_types.rs
new file mode 100644
index 00000000..134a8966
--- /dev/null
+++ b/bindings/python/src/merge_types.rs
@@ -0,0 +1,228 @@
+// 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.
+
+//! Immutable Python descriptions of core MERGE clauses and assignments.
+
+use pyo3::exceptions::PyTypeError;
+use pyo3::prelude::*;
+use pyo3::types::PyMapping;
+use pyo3::{PyTraverseError, PyVisit};
+
+pub(crate) enum AssignmentValue {
+ SourceColumn(String),
+ TargetColumn(String),
+ Literal(Py<PyAny>),
+ Function(Py<PyAny>),
+}
+
+/// A typed column reference, literal, or function evaluated for selected rows.
+#[pyclass(name = "MergeAssignment", module = "pypaimon_rust.datafusion",
frozen)]
+pub struct PyMergeAssignment {
+ pub(crate) value: AssignmentValue,
+}
+
+#[pymethods]
+impl PyMergeAssignment {
+ #[staticmethod]
+ fn source_column(name: String) -> Self {
+ Self {
+ value: AssignmentValue::SourceColumn(name),
+ }
+ }
+
+ #[staticmethod]
+ fn target_column(name: String) -> Self {
+ Self {
+ value: AssignmentValue::TargetColumn(name),
+ }
+ }
+
+ #[staticmethod]
+ fn literal(value: Py<PyAny>) -> Self {
+ Self {
+ value: AssignmentValue::Literal(value),
+ }
+ }
+
+ #[staticmethod]
+ fn function(function: Py<PyAny>, py: Python<'_>) -> PyResult<Self> {
+ if !function.bind(py).is_callable() {
+ return Err(PyTypeError::new_err(
+ "MergeAssignment.function requires a callable",
+ ));
+ }
+ Ok(Self {
+ value: AssignmentValue::Function(function),
+ })
+ }
+
+ fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+ match &self.value {
+ AssignmentValue::Literal(value) | AssignmentValue::Function(value)
=> visit.call(value),
+ AssignmentValue::SourceColumn(_) |
AssignmentValue::TargetColumn(_) => Ok(()),
+ }
+ }
+}
+
+pub(crate) type Assignments = Vec<(String, Py<PyMergeAssignment>)>;
+
+fn assignments(values: &Bound<'_, PyMapping>) -> PyResult<Assignments> {
+ values.items()?.extract()
+}
+
+/// An ordered WHEN MATCHED UPDATE or DELETE branch.
+#[pyclass(name = "WhenMatched", module = "pypaimon_rust.datafusion", frozen)]
+pub struct PyWhenMatched {
+ pub(crate) condition: Option<String>,
+ pub(crate) delete: bool,
+ pub(crate) assignments: Assignments,
+}
+
+#[pymethods]
+impl PyWhenMatched {
+ #[staticmethod]
+ #[pyo3(signature = (assignments, condition=None))]
+ fn update(assignments: &Bound<'_, PyMapping>, condition: Option<String>)
-> PyResult<Self> {
+ Ok(Self {
+ condition,
+ delete: false,
+ assignments: self::assignments(assignments)?,
+ })
+ }
+
+ #[staticmethod]
+ #[pyo3(signature = (condition=None))]
+ fn delete(condition: Option<String>) -> Self {
+ Self {
+ condition,
+ delete: true,
+ assignments: Vec::new(),
+ }
+ }
+
+ fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+ for (_, value) in &self.assignments {
+ visit.call(value)?;
+ }
+ Ok(())
+ }
+}
+
+/// An ordered WHEN NOT MATCHED INSERT branch.
+#[pyclass(name = "WhenNotMatched", module = "pypaimon_rust.datafusion",
frozen)]
+pub struct PyWhenNotMatched {
+ pub(crate) condition: Option<String>,
+ pub(crate) assignments: Assignments,
+}
+
+#[pymethods]
+impl PyWhenNotMatched {
+ #[new]
+ #[pyo3(signature = (assignments, condition=None))]
+ fn new(assignments: &Bound<'_, PyMapping>, condition: Option<String>) ->
PyResult<Self> {
+ Ok(Self {
+ condition,
+ assignments: self::assignments(assignments)?,
+ })
+ }
+
+ fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+ for (_, value) in &self.assignments {
+ visit.call(value)?;
+ }
+ Ok(())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use pyo3::ffi::c_str;
+
+ #[test]
+ fn typed_factories_validate_descriptions_and_snapshot_assignments() {
+ Python::attach(|py| {
+ let locals = pyo3::types::PyDict::new(py);
+ locals
+ .set_item("MergeAssignment",
py.get_type::<PyMergeAssignment>())
+ .unwrap();
+ locals
+ .set_item("WhenMatched", py.get_type::<PyWhenMatched>())
+ .unwrap();
+ locals
+ .set_item("WhenNotMatched", py.get_type::<PyWhenNotMatched>())
+ .unwrap();
+ py.run(
+ c_str!(
+ r#"
+from types import MappingProxyType
+value = MergeAssignment.source_column('value')
+mapping = {'value': value}
+matched = WhenMatched.update(MappingProxyType(mapping), condition='"s.value" >
0')
+unmatched = WhenNotMatched(mapping)
+mapping.clear()
+for make in [
+ lambda: MergeAssignment(),
+ lambda: MergeAssignment.source_column(1),
+ lambda: MergeAssignment.target_column(None),
+ lambda: MergeAssignment.function(1),
+ lambda: WhenMatched(),
+ lambda: WhenMatched.update({'value': 1}),
+ lambda: WhenMatched.update({1: value}),
+ lambda: WhenMatched.update({'value': value}, condition={'sql': 'true'}),
+ lambda: WhenMatched.delete(condition={'sql': 'true'}),
+ lambda: WhenNotMatched({'value': 1}),
+ lambda: WhenNotMatched({'value': value}, condition={'sql': 'true'}),
+]:
+ try:
+ make()
+ except TypeError:
+ pass
+ else:
+ raise AssertionError('untyped MERGE description was accepted')
+for description in [value, matched, unmatched]:
+ try:
+ description.condition = None
+ except AttributeError:
+ pass
+ else:
+ raise AssertionError('MERGE description was mutable')
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ let matched = locals
+ .get_item("matched")
+ .unwrap()
+ .unwrap()
+ .extract::<PyRef<PyWhenMatched>>()
+ .unwrap();
+ let unmatched = locals
+ .get_item("unmatched")
+ .unwrap()
+ .unwrap()
+ .extract::<PyRef<PyWhenNotMatched>>()
+ .unwrap();
+ assert_eq!(matched.assignments.len(), 1);
+ assert_eq!(unmatched.assignments.len(), 1);
+ assert_eq!(matched.assignments[0].0, "value");
+ assert_eq!(unmatched.assignments[0].0, "value");
+ });
+ }
+}
diff --git a/bindings/python/src/update_assignment.rs
b/bindings/python/src/update_assignment.rs
index 3df5a3a5..0779cbb7 100644
--- a/bindings/python/src/update_assignment.rs
+++ b/bindings/python/src/update_assignment.rs
@@ -72,72 +72,85 @@ pub(crate) fn from_python_schema(
schema: SchemaRef,
error: CallbackError,
) -> PyResult<Vec<(String, UpdateAssignment)>> {
- let pa = py.import("pyarrow")?;
- let array_type = pa.getattr("Array")?;
- let chunked_type = pa.getattr("ChunkedArray")?;
assignments
.iter()
.map(|(name, value)| {
let name: String = name.extract()?;
let assignment = if value.is_callable() {
- let value = value.unbind();
- let error = error.clone();
- UpdateAssignment::Function(Arc::new(move |batches| {
- // A Python callback may invoke another native read or
write.
- tokio::task::block_in_place(|| {
- Python::attach(|py| -> PyResult<_> {
- let pa = py.import("pyarrow")?;
- let batches = batches
- .iter()
- .map(|batch| batch.to_pyarrow(py))
- .collect::<PyResult<Vec<_>>>()?;
- let matched = pa
- .getattr("Table")?
- .call_method1("from_batches", (PyList::new(py,
batches)?,))?;
- arrow_arrays(py,
&value.bind(py).call1((matched,))?)
- })
- })
- .map_err(|cause| bridge_error(cause, &error))
- }))
- } else if value.is_instance(&array_type)? ||
value.is_instance(&chunked_type)? {
- UpdateAssignment::Array(arrow_arrays(py, &value)?)
+ function(value.unbind(), error.clone())
} else {
- // Literal conversion stays lazy: a no-match update must
neither
- // evaluate callbacks nor reject a scalar that is never
assigned.
- let value = value.unbind();
- let name = name.clone();
- let schema = schema.clone();
- let error = error.clone();
- UpdateAssignment::DeferredScalar(Arc::new(move || {
- tokio::task::block_in_place(|| {
- Python::attach(|py| -> PyResult<_> {
- let pa = py.import("pyarrow")?;
- let value = value.bind(py);
- let scalar = if
value.is_instance(&pa.getattr("Scalar")?)? {
- value.call_method0("as_py")?
- } else {
- value.clone()
- };
- let kwargs = PyDict::new(py);
- kwargs.set_item(
- "type",
- schema
- .to_pyarrow(py)?
- .call_method1("field", (&name,))?
- .getattr("type")?,
- )?;
- let one = pa.call_method(
- "array",
- (PyList::new(py, [scalar])?,),
- Some(&kwargs),
- )?;
- crate::arrow_input::array_from_pyarrow(&one)
- })
- })
- .map_err(|cause| bridge_error(cause, &error))
- }))
+ literal(
+ py,
+ name.clone(),
+ value.unbind(),
+ schema.clone(),
+ error.clone(),
+ )?
};
Ok((name, assignment))
})
.collect()
}
+
+pub(crate) fn function(value: Py<PyAny>, error: CallbackError) ->
UpdateAssignment {
+ UpdateAssignment::Function(Arc::new(move |batches| {
+ // A Python callback may invoke another native read or write.
+ tokio::task::block_in_place(|| {
+ Python::attach(|py| -> PyResult<_> {
+ let pa = py.import("pyarrow")?;
+ let batches = batches
+ .iter()
+ .map(|batch| batch.to_pyarrow(py))
+ .collect::<PyResult<Vec<_>>>()?;
+ let matched = pa
+ .getattr("Table")?
+ .call_method1("from_batches", (PyList::new(py,
batches)?,))?;
+ arrow_arrays(py, &value.bind(py).call1((matched,))?)
+ })
+ })
+ .map_err(|cause| bridge_error(cause, &error))
+ }))
+}
+
+pub(crate) fn literal(
+ py: Python<'_>,
+ name: String,
+ value: Py<PyAny>,
+ schema: SchemaRef,
+ error: CallbackError,
+) -> PyResult<UpdateAssignment> {
+ let pa = py.import("pyarrow")?;
+ let bound = value.bind(py);
+ if bound.is_instance(&pa.getattr("Array")?)?
+ || bound.is_instance(&pa.getattr("ChunkedArray")?)?
+ {
+ return Ok(UpdateAssignment::Array(arrow_arrays(py, bound)?));
+ }
+ // Literal conversion stays lazy: a no-match update must neither evaluate
+ // callbacks nor reject a scalar that is never assigned. Explicit MERGE
+ // literals also stay literals when the Python value happens to be
callable.
+ Ok(UpdateAssignment::DeferredScalar(Arc::new(move || {
+ tokio::task::block_in_place(|| {
+ Python::attach(|py| -> PyResult<_> {
+ let pa = py.import("pyarrow")?;
+ let value = value.bind(py);
+ let scalar = if value.is_instance(&pa.getattr("Scalar")?)? {
+ value.call_method0("as_py")?
+ } else {
+ value.clone()
+ };
+ let kwargs = PyDict::new(py);
+ kwargs.set_item(
+ "type",
+ schema
+ .to_pyarrow(py)?
+ .call_method1("field", (&name,))?
+ .getattr("type")?,
+ )?;
+ let one = pa.call_method("array", (PyList::new(py,
[scalar])?,), Some(&kwargs))?;
+ crate::arrow_input::array_from_pyarrow(&one)
+ })
+ })
+ .map_err(|cause| bridge_error(cause, &error))
+ })))
+}
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index e5688b9b..af66d51e 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -1847,6 +1847,29 @@ mod tests {
use paimon::spec::{BigIntType, IntType, Schema, TableSchema, VarCharType};
use pyo3::ffi::c_str;
+ fn merge_locals(py: Python<'_>) -> Bound<'_, PyDict> {
+ let locals = PyDict::new(py);
+ locals
+ .set_item(
+ "MergeAssignment",
+ py.get_type::<crate::merge_types::PyMergeAssignment>(),
+ )
+ .unwrap();
+ locals
+ .set_item(
+ "WhenMatched",
+ py.get_type::<crate::merge_types::PyWhenMatched>(),
+ )
+ .unwrap();
+ locals
+ .set_item(
+ "WhenNotMatched",
+ py.get_type::<crate::merge_types::PyWhenNotMatched>(),
+ )
+ .unwrap();
+ locals
+ }
+
#[test]
fn
captured_plan_rejects_other_operations_before_python_input_consumption() {
Python::attach(|py| {
@@ -2406,6 +2429,174 @@ first = install('first')
});
}
+ #[test]
+ fn merge_into_rejects_legacy_dictionary_clauses() {
+ Python::attach(|py| {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("row-tracking.enabled", "true")
+ .option("data-evolution.enabled", "true")
+ .build()
+ .unwrap();
+ let table = Arc::new(Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "typed_merge_clauses"),
+ "memory:/typed-merge-clauses".into(),
+ TableSchema::new(0, &schema),
+ None,
+ ));
+ let update = Py::new(py, PyBatchWriteBuilder::new(table))
+ .unwrap()
+ .bind(py)
+ .call_method0("new_update")
+ .unwrap();
+ let locals = PyDict::new(py);
+ locals.set_item("update", update).unwrap();
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+source = pa.table({'id': pa.array([], pa.int32())})
+for matched, unmatched in [
+ ([dict(condition=None, delete=False, assignments=[('id', 'source',
'id')])], []),
+ ([], [dict(condition=None, assignments=[('id', 'source', 'id')])]),
+]:
+ try:
+ update.merge_into(source, on=[('id', 'id')],
+ when_matched=matched, when_not_matched=unmatched)
+ except TypeError:
+ pass
+ else:
+ raise AssertionError('MERGE accepted a legacy dictionary clause')
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn typed_merge_into_preserves_literal_and_callback_evaluation() {
+ use arrow::array::Int32Array;
+ use futures::TryStreamExt;
+
+ Python::attach(|py| {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .option("row-tracking.enabled", "true")
+ .option("data-evolution.enabled", "true")
+ .build()
+ .unwrap();
+ let table = Arc::new(Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "typed_merge_lazy_values"),
+ "memory:/typed-merge-lazy-values".into(),
+ TableSchema::new(0, &schema),
+ None,
+ ));
+ let locals = merge_locals(py);
+ locals
+ .set_item(
+ "builder",
+ Py::new(py,
PyBatchWriteBuilder::new(table.clone())).unwrap(),
+ )
+ .unwrap();
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+source = pa.table({'id': pa.array([1, 2], pa.int32()),
+ 'value': pa.array([10, 20], pa.int32())})
+update = builder.new_update()
+literal_calls = []
+def literal_value(rows):
+ literal_calls.append(rows.num_rows)
+ return pa.array([99, 99], pa.int32())
+unused = WhenNotMatched({'id': MergeAssignment.source_column('id'),
+ 'value': MergeAssignment.literal(literal_value)},
+ condition='"s.id" < 0')
+assert update.merge_into(source, on=[('id', 'id')], when_matched=[],
+ when_not_matched=[unused]) == []
+assert literal_calls == []
+selected = WhenNotMatched({'id': MergeAssignment.source_column('id'),
+ 'value': MergeAssignment.literal(literal_value)})
+try:
+ update.merge_into(source, on=[('id', 'id')], when_matched=[],
+ when_not_matched=[selected])
+except (TypeError, ValueError):
+ pass
+else:
+ raise AssertionError('callable literal was evaluated as an assignment
function')
+assert literal_calls == []
+
+failure = LookupError('typed MERGE callback failed')
+callback_rows = []
+def assignment(rows):
+ callback_rows.append(rows['s.id'].to_pylist())
+ if len(callback_rows) == 1:
+ raise failure
+ return pa.chunked_array([pa.array([21], pa.int32())])
+assignments = {'id': MergeAssignment.source_column('id'),
+ 'value': MergeAssignment.function(function=assignment)}
+clause = WhenNotMatched(assignments, condition='"s.id" = 2')
+assignments.clear()
+for attempt in [1, 2]:
+ try:
+ messages = update.merge_into(source, on=[('id', 'id')],
when_matched=[],
+ when_not_matched=[unused, clause])
+ except LookupError as error:
+ assert attempt == 1
+ assert error is failure
+ else:
+ assert attempt == 2
+ builder.new_commit().commit(messages)
+assert callback_rows == [[2], [2]]
+assert literal_calls == []
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ let batches = py.detach(|| {
+ runtime().block_on(async {
+ let builder = table.new_read_builder();
+ let plan = builder.new_scan().plan().await.unwrap();
+ builder
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap()
+ })
+ });
+
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
+ assert_eq!(
+ batches[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 2
+ );
+ assert_eq!(
+ batches[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 21
+ );
+ });
+ }
+
#[test]
fn merge_table_source_conditions_resolve_read_metadata() {
Python::attach(|py| {
@@ -2434,15 +2625,15 @@ first = install('first')
));
let source = Py::new(py,
crate::table::PyTable::new(source)).unwrap();
let context =
UpdateContext::new(&WriteContext::new(table)).unwrap();
- let locals = PyDict::new(py);
+ let locals = merge_locals(py);
locals.set_item("metadata", metadata).unwrap();
py.run(
c_str!(
r#"
-condition = dict(sql='"s.' + metadata + '" >= 0')
-matched = [dict(assignments=[('value', 'source', 'value')],
condition=condition, delete=False)]
-not_matched = [dict(assignments=[('id', 'source', metadata), ('value',
'source', 'value')],
- condition=condition)]
+condition = '"s.' + metadata + '" >= 0'
+matched = [WhenMatched.update({'value':
MergeAssignment.source_column('value')}, condition=condition)]
+not_matched = [WhenNotMatched({'id': MergeAssignment.source_column(metadata),
+ 'value':
MergeAssignment.source_column('value')}, condition=condition)]
"#
),
Some(&locals),
@@ -2558,17 +2749,18 @@ not_matched = [dict(assignments=[('id', 'source',
metadata), ('value', 'source',
});
let source = Py::new(py,
crate::table::PyTable::new(source)).unwrap();
let context =
UpdateContext::new(&WriteContext::new(target.clone())).unwrap();
- let locals = PyDict::new(py);
+ let locals = merge_locals(py);
locals.set_item("metadata", metadata).unwrap();
locals.set_item("first", first).unwrap();
locals.set_item("second", second).unwrap();
py.run(
c_str!(
r#"
-matched = [dict(assignments=[('value', 'source', 'value')],
- condition=dict(sql='"s.' + metadata + '" = ' + str(first)),
delete=False)]
-not_matched = [dict(assignments=[('id', 'source', metadata), ('value',
'source', 'value')],
- condition=dict(sql='"s.' + metadata + '" = ' +
str(second)))]
+matched = [WhenMatched.update({'value':
MergeAssignment.source_column('value')},
+ condition='"s.' + metadata + '" = ' +
str(first))]
+not_matched = [WhenNotMatched({'id': MergeAssignment.source_column(metadata),
+ 'value': MergeAssignment.source_column('value')},
+ condition='"s.' + metadata + '" = ' +
str(second))]
"#
),
Some(&locals),
@@ -2794,7 +2986,7 @@ invalid = source.slice(0, 1)
})
});
let context =
UpdateContext::new(&WriteContext::new(table)).unwrap();
- let locals = PyDict::new(py);
+ let locals = merge_locals(py);
locals.set_item("reverse", reverse).unwrap();
py.run(
c_str!(
@@ -2817,7 +3009,7 @@ encoded =
pa.chunked_array([pa.DictionaryArray.from_arrays(
source = source.append_column(pa.field('encoded_unused', encoded.type,
nullable=False), encoded)
assert source.column('encoded_unused').null_count == 0
assert source.column('encoded_unused').to_pylist() == [None, None, None]
-matched = [dict(assignments=[('value', 'source', 'value')], condition=None,
delete=False)]
+matched = [WhenMatched.update({'value':
MergeAssignment.source_column('value')})]
not_matched = []
"#
),
diff --git a/bindings/python/tests/test_blob_uri_reader.py
b/bindings/python/tests/test_blob_uri_reader.py
index 54bbc54a..aa457e22 100644
--- a/bindings/python/tests/test_blob_uri_reader.py
+++ b/bindings/python/tests/test_blob_uri_reader.py
@@ -22,7 +22,7 @@ import pyarrow as pa
import pytest
from test_blob_consumer import _table
-from pypaimon_rust.datafusion import PaimonCatalog, SQLContext
+from pypaimon_rust.datafusion import MergeAssignment, PaimonCatalog,
SQLContext, WhenMatched
def descriptor(offset, length):
@@ -217,9 +217,9 @@ def
test_update_factory_errors_keep_identity_and_existing_data(tmp_path, stream,
extra = {'commit_identifier': 7} if stream else {}
with pytest.raises(OSError) as caught:
if operation == 'merge':
- update.merge_into(source, on=[('id', 'id')], when_matched=[{
- 'delete': False, 'assignments': [('payload', 'source',
'payload')],
- }], when_not_matched=[], **extra)
+ update.merge_into(source, on=[('id', 'id')], when_matched=[
+ WhenMatched.update({'payload':
MergeAssignment.source_column('payload')}),
+ ], when_not_matched=[], **extra)
elif operation == 'row_id':
update.update_by_arrow_with_row_id(data, **extra)
elif operation == 'grouped_row_id':
diff --git a/bindings/python/tests/test_update.py
b/bindings/python/tests/test_update.py
index 069a3a34..83dc9464 100644
--- a/bindings/python/tests/test_update.py
+++ b/bindings/python/tests/test_update.py
@@ -19,7 +19,9 @@ import pyarrow as pa
import pytest
import pypaimon_rust.datafusion as datafusion
-from pypaimon_rust.datafusion import PaimonCatalog, SQLContext
+from pypaimon_rust.datafusion import (
+ MergeAssignment, PaimonCatalog, SQLContext, WhenMatched, WhenNotMatched,
+)
def
test_upsert_batches_preserve_missing_fields_nulls_and_stream_updates(tmp_path):
@@ -449,13 +451,14 @@ def
test_prepared_merge_conditions_chunks_and_lazy_literals(tmp_path, stream):
('id', pa.int32()), ('value', pa.int32())]))
update = builder.new_update()
kwargs = dict(commit_identifier=42) if stream else {}
- false_clause = dict(delete=False, condition=dict(sql='"s.value" < 0'),
- assignments=[('value', 'literal', 'invalid int')])
+ false_clause = WhenMatched.update(
+ {'value': MergeAssignment.literal('invalid int')},
condition='"s.value" < 0')
assert update.merge_into(source, on=[('id', 'id')],
when_matched=[false_clause],
when_not_matched=[], **kwargs) == []
messages = update.merge_into(source, on=[('id', 'id')], when_matched=[
- dict(delete=False, condition=dict(sql='"s.value" IN (SELECT v FROM
(VALUES (11), (22)) AS m(v))'),
- assignments=[('value', 'literal', pa.chunked_array([[100],
[200]]))])],
+ WhenMatched.update(
+ {'value': MergeAssignment.literal(pa.chunked_array([[100],
[200]]))},
+ condition='"s.value" IN (SELECT v FROM (VALUES (11), (22)) AS
m(v))')],
when_not_matched=[], **kwargs)
# Preparation leaves the visible snapshot unchanged.
@@ -495,12 +498,14 @@ def
test_prepared_merge_accepts_native_table_sources(tmp_path, stream, historica
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')
+ condition = '"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)
+ WhenMatched.update({'value':
MergeAssignment.source_column('amount')}, condition)],
+ when_not_matched=[WhenNotMatched({
+ 'id': MergeAssignment.source_column('source_id'),
+ 'value': MergeAssignment.source_column('amount'),
+ }, condition)], **kwargs)
assert target.latest_snapshot().id() == target_snapshot
assert source.latest_snapshot().id() == source_snapshot
commit = builder.new_commit()
@@ -533,8 +538,9 @@ def
test_prepared_merge_callable_uses_complete_clause_and_preserves_exception(tm
return pa.chunked_array([[100], [200]], type=pa.int32())
messages = update.merge_into(source, on=[('id', 'id')], when_matched=[
- dict(delete=False, assignments=[('value', 'callable', values)])],
when_not_matched=[
- dict(assignments=[('id', 'source', 'id'), ('value', 'callable', lambda
batch: batch['s.delta'])])],
+ WhenMatched.update({'value': MergeAssignment.function(values)})],
when_not_matched=[
+ WhenNotMatched({'id': MergeAssignment.source_column('id'),
+ 'value': MergeAssignment.function(lambda batch:
batch['s.delta'])})],
**kwargs)
assert calls == [2]
commit = builder.new_commit()
@@ -547,7 +553,7 @@ def
test_prepared_merge_callable_uses_complete_clause_and_preserves_exception(tm
def fail(batch):
raise error
- clause = [dict(delete=False, assignments=[('value', 'callable', fail)])]
+ clause = [WhenMatched.update({'value': MergeAssignment.function(fail)})]
update = builder.new_update()
assert update.merge_into(pa.table({'id': pa.array([99], pa.int32())}),
on=[('id', 'id')],
when_matched=clause, when_not_matched=[],
**kwargs) == []
