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 9f8338d9 Support Arrow row ID inputs for deletion (#1210)
9f8338d9 is described below
commit 9f8338d91a1c3d5a5ff07f0a7716f7f7f256f518
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 21:43:30 2026 +0800
Support Arrow row ID inputs for deletion (#1210)
---
bindings/python/README.md | 4 +
.../python/python/pypaimon_rust/datafusion.pyi | 11 +-
bindings/python/src/write.rs | 175 ++++++++++++++++++++-
bindings/python/tests/test_update.py | 145 +++++++++++++++++
crates/paimon/src/table/table_update.rs | 8 +-
crates/paimon/src/table/table_update_plan_tests.rs | 22 ++-
.../tests/data_evolution_row_sidecar_test.rs | 2 +-
7 files changed, 354 insertions(+), 13 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index d793d2d3..696e746a 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -227,6 +227,10 @@ authorization and the captured snapshot, and adds sorted,
deduplicated row range
for reading. It does not check row liveness or authorize file access; readers
and mutation workers retain those checks. The selected plan is only for
reading.
+Batch and stream updates accept integer Arrow arrays and chunked arrays through
+`delete_by_row_id`, as well as integer sequences. Arrow chunks and slices are
+imported directly; nulls and values outside Int64 are rejected before writing.
+
### Shard column updates
The batch and stream update objects expose
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 92964a99..0ede5a1f 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -638,7 +638,14 @@ class BatchTableUpdate:
"""
...
def with_update_type(self, update_cols: List[str]) -> "BatchTableUpdate":
...
- def delete_by_row_id(self, row_ids: Sequence[int]) -> List[CommitMessage]:
...
+ def delete_by_row_id(
+ self, row_ids: Union[Sequence[int], pyarrow.Array,
pyarrow.ChunkedArray]
+ ) -> List[CommitMessage]:
+ """Prepare row deletes from integer IDs; Arrow inputs retain chunks
and slices.
+
+ Nulls and values outside Int64 are rejected. Duplicate IDs are deleted
once.
+ """
+ ...
def delete_by_predicate(
self, predicate: Optional[Dict[str, Any]]
) -> List[CommitMessage]:
@@ -712,7 +719,7 @@ class StreamTableUpdate:
...
def with_update_type(self, update_cols: List[str]) -> "StreamTableUpdate":
...
def delete_by_row_id(
- self, row_ids: Sequence[int], commit_identifier: int
+ self, row_ids: Union[Sequence[int], pyarrow.Array,
pyarrow.ChunkedArray], commit_identifier: int
) -> List[CommitMessage]: ...
def delete_by_predicate(
self, predicate: Optional[Dict[str, Any]], commit_identifier: int
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index 7a8b0b91..0e1009c4 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -18,7 +18,9 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
-use arrow::datatypes::Schema as ArrowSchema;
+use arrow::array::{Array, Int64Array};
+use arrow::compute::{cast_with_options, CastOptions};
+use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema};
use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use arrow::record_batch::RecordBatch;
use paimon::spec::{BlobConsumer, BlobDescriptor, CoreOptions, DataType, Datum,
RowType};
@@ -28,7 +30,7 @@ use paimon::table::{
COMMIT_MESSAGE_SERIALIZER_VERSION,
};
use paimon_datafusion::runtime::runtime;
-use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
+use pyo3::exceptions::{PyModuleNotFoundError, PyRuntimeError, PyTypeError,
PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyString};
use pyo3::{PyTraverseError, PyVisit};
@@ -41,6 +43,75 @@ use crate::predicate::{dict_to_table_predicate, py_to_datum};
type PartitionSpec = HashMap<String, Option<Datum>>;
type PythonPartitionSpec = HashMap<String, Py<PyAny>>;
+fn delete_row_id_arrays(row_ids: &Bound<'_, PyAny>) ->
PyResult<Vec<Int64Array>> {
+ let py = row_ids.py();
+ let pa = match py.import("pyarrow") {
+ Ok(pa) => pa,
+ // Plain sequence deletes do not require the optional PyArrow package.
+ // Preserve errors from a broken installation or its dependencies.
+ Err(error)
+ if error.is_instance_of::<PyModuleNotFoundError>(py)
+ && error
+ .value(py)
+ .getattr("name")?
+ .extract::<Option<String>>()?
+ .as_deref()
+ == Some("pyarrow") =>
+ {
+ return Ok(vec![Int64Array::from(row_ids.extract::<Vec<i64>>()?)]);
+ }
+ Err(error) => return Err(error),
+ };
+ let is_array = row_ids.is_instance(&pa.getattr("Array")?)?;
+ let is_chunked = row_ids.is_instance(&pa.getattr("ChunkedArray")?)?;
+ if !is_array && !is_chunked {
+ return Ok(vec![Int64Array::from(row_ids.extract::<Vec<i64>>()?)]);
+ }
+ // Check the declared type even when a ChunkedArray contains no chunks.
+ if
!ArrowDataType::from_pyarrow_bound(&row_ids.getattr("type")?)?.is_integer() {
+ return Err(PyTypeError::new_err(
+ "_ROW_ID input must be an integer Arrow array",
+ ));
+ }
+ let arrays = if is_array {
+ vec![crate::arrow_input::array_from_pyarrow(row_ids)?]
+ } else {
+ row_ids
+ .getattr("chunks")?
+ .try_iter()?
+ .map(|chunk| crate::arrow_input::array_from_pyarrow(&chunk?))
+ .collect::<PyResult<Vec<_>>>()?
+ };
+ py.detach(|| {
+ arrays
+ .into_iter()
+ .map(|array| {
+ if array.null_count() != 0 {
+ return Err(PyValueError::new_err("_ROW_ID must not be
null"));
+ }
+ let converted = cast_with_options(
+ array.as_ref(),
+ &ArrowDataType::Int64,
+ &CastOptions {
+ safe: false,
+ ..Default::default()
+ },
+ )
+ .map_err(|error| {
+ PyValueError::new_err(format!(
+ "_ROW_ID cannot be represented as Int64: {error}"
+ ))
+ })?;
+ Ok(converted
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .expect("cast to Int64 returned an Int64Array")
+ .clone())
+ })
+ .collect()
+ })
+}
+
/// Shared implementation state; public builders keep batch and stream
contracts separate.
struct WriteContext {
table: Arc<Table>,
@@ -880,10 +951,19 @@ impl UpdateContext {
fn delete_by_row_id(
&self,
py: Python<'_>,
- row_ids: Vec<i64>,
+ row_ids: &Bound<'_, PyAny>,
) -> PyResult<Vec<PyCommitMessage>> {
+ let arrays = delete_row_id_arrays(row_ids)?;
let messages = py
- .detach(||
runtime().block_on(self.inner.delete_by_row_id(row_ids)))
+ .detach(|| {
+ runtime().block_on(
+ self.inner.delete_by_row_id(
+ arrays
+ .iter()
+ .flat_map(|array| array.values().iter().copied()),
+ ),
+ )
+ })
.map_err(to_py_err)?;
Ok(wrap_messages(
messages,
@@ -1193,7 +1273,7 @@ impl PyStreamTableUpdate {
fn delete_by_row_id(
&self,
py: Python<'_>,
- row_ids: Vec<i64>,
+ row_ids: &Bound<'_, PyAny>,
commit_identifier: i64,
) -> PyResult<Vec<PyCommitMessage>> {
// The stream committer applies the identifier to the returned
messages.
@@ -1327,7 +1407,7 @@ impl PyBatchTableUpdate {
fn delete_by_row_id(
&self,
py: Python<'_>,
- row_ids: Vec<i64>,
+ row_ids: &Bound<'_, PyAny>,
) -> PyResult<Vec<PyCommitMessage>> {
self.context.delete_by_row_id(py, row_ids)
}
@@ -1871,6 +1951,89 @@ mod tests {
locals
}
+ #[test]
+ fn row_id_delete_inputs_preserve_arrow_chunks_and_integer_values() {
+ Python::attach(|py| {
+ let locals = PyDict::new(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+valid = [
+ ([2, 1], [2, 1], [2]),
+ ((2, 1), [2, 1], [2]),
+ (pa.array([None, 2, 1, None], type=pa.int64()).slice(1, 2), [2, 1], [2]),
+ (pa.chunked_array([[999, 2], [1, 999]], type=pa.int16()).slice(1, 2), [2,
1], [1, 1]),
+ (pa.array([2 ** 63 - 1], type=pa.uint64()), [2 ** 63 - 1], [1]),
+ (pa.array([], type=pa.uint32()), [], [0]),
+ (pa.chunked_array([], type=pa.int64()), [], []),
+]
+invalid = [
+ (pa.array([None], type=pa.int64()), 'value'),
+ (pa.chunked_array([[0], [None]], type=pa.int64()), 'value'),
+ (pa.array([2 ** 63], type=pa.uint64()), 'value'),
+ (pa.chunked_array([[0], [2 ** 63]], type=pa.uint64()), 'value'),
+ (pa.array([1.5]), 'type'),
+ (pa.array([True]), 'type'),
+ (pa.array([], type=pa.bool_()), 'type'),
+ (pa.array([], type=pa.string()), 'type'),
+ (pa.chunked_array([], type=pa.float64()), 'type'),
+ ([1.5], 'type'),
+ ([None], 'type'),
+ ([2 ** 63], 'overflow'),
+]
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ for case in locals
+ .get_item("valid")
+ .unwrap()
+ .unwrap()
+ .try_iter()
+ .unwrap()
+ {
+ let case = case.unwrap();
+ let arrays =
delete_row_id_arrays(&case.get_item(0).unwrap()).unwrap();
+ let values: Vec<i64> = arrays
+ .iter()
+ .flat_map(|array| array.values().iter().copied())
+ .collect();
+ let lengths: Vec<usize> =
arrays.iter().map(Int64Array::len).collect();
+ assert_eq!(
+ values,
+ case.get_item(1).unwrap().extract::<Vec<i64>>().unwrap()
+ );
+ assert_eq!(
+ lengths,
+ case.get_item(2).unwrap().extract::<Vec<usize>>().unwrap()
+ );
+ }
+ for case in locals
+ .get_item("invalid")
+ .unwrap()
+ .unwrap()
+ .try_iter()
+ .unwrap()
+ {
+ let case = case.unwrap();
+ let error =
delete_row_id_arrays(&case.get_item(0).unwrap()).unwrap_err();
+ let kind: String =
case.get_item(1).unwrap().extract().unwrap();
+ assert!(
+ match kind.as_str() {
+ "type" => error.is_instance_of::<PyTypeError>(py),
+ "value" => error.is_instance_of::<PyValueError>(py),
+ "overflow" =>
error.is_instance_of::<pyo3::exceptions::PyOverflowError>(py),
+ _ => false,
+ },
+ "unexpected {kind} error: {error}"
+ );
+ }
+ });
+ }
+
#[test]
fn
captured_plan_rejects_other_operations_before_python_input_consumption() {
Python::attach(|py| {
diff --git a/bindings/python/tests/test_update.py
b/bindings/python/tests/test_update.py
index 83dc9464..84f8be27 100644
--- a/bindings/python/tests/test_update.py
+++ b/bindings/python/tests/test_update.py
@@ -24,6 +24,151 @@ from pypaimon_rust.datafusion import (
)
[email protected]('stream', [False, True])
+def test_arrow_row_id_deletes_preserve_captured_bucket_vectors(tmp_path,
stream):
+ context = SQLContext()
+ context.register_catalog('paimon', {'warehouse': str(tmp_path)})
+ context.sql('CREATE SCHEMA paimon.arrow_deletes')
+ context.sql("""CREATE TABLE paimon.arrow_deletes.t (id INT) WITH (
+ 'bucket' = '-1', 'row-tracking.enabled' = 'true',
+ 'data-evolution.enabled' = 'true', 'deletion-vectors.enabled' =
'true')""")
+ for values in ['(10), (11)', '(20), (21)']:
+ context.sql('INSERT INTO paimon.arrow_deletes.t (id) VALUES ' + values)
+ table = PaimonCatalog({'warehouse':
str(tmp_path)}).get_table('arrow_deletes.t')
+ builder =
(table.new_stream_write_builder().with_commit_user('arrow-deletes')
+ if stream else table.new_batch_write_builder())
+
+ def delete(update, row_ids, identifier):
+ return (update.delete_by_row_id(row_ids, identifier) if stream
+ else update.delete_by_row_id(row_ids))
+
+ def commit(messages, identifier):
+ committer = builder.new_commit()
+ committer.commit(identifier, messages) if stream else
committer.commit(messages)
+
+ def rows():
+ return pa.Table.from_batches(context.sql(
+ 'SELECT id FROM paimon.arrow_deletes.t ORDER BY
id')).column('id').to_pylist()
+
+ # Keep the legacy sequence input and seed a DV for a different anchor.
+ commit(delete(builder.new_update(), [0], 40), 40)
+ plan = table.new_read_builder().new_scan().plan_all()
+ groups = datafusion.RowIdPlan(plan).groups()
+ selected = datafusion.Plan(plan.snapshot_id(), [split for first, _, split
in groups if first >= 2])
+ update = builder.new_update().with_plan(selected)
+ context.sql('INSERT INTO paimon.arrow_deletes.t (id) VALUES (30)')
+
+ ids = pa.chunked_array([
+ pa.array([999, 2, 2], type=pa.uint64()).slice(1),
+ pa.array([3, 999], type=pa.uint64()).slice(0, 1),
+ ])
+ messages = delete(update, ids, 41)
+ stale = delete(update, pa.array([3], type=pa.int16()), 42)
+ assert rows() == [11, 20, 21, 30]
+ commit(messages, 41)
+ assert rows() == [11, 30]
+ # A second preparation from the captured snapshot cannot replace its old
+ # index again after the first publication.
+ snapshot = table.latest_snapshot().id()
+ with pytest.raises(ValueError, match='missing deletion vector index'):
+ commit(stale, 42)
+ assert table.latest_snapshot().id() == snapshot
+ assert rows() == [11, 30]
+
+
[email protected]('stream', [False, True])
+def test_arrow_row_id_delete_input_validation(tmp_path, stream):
+ from pathlib import Path
+
+ context = SQLContext()
+ context.register_catalog('paimon', {'warehouse': str(tmp_path)})
+ context.sql('CREATE SCHEMA paimon.delete_inputs')
+ context.sql("""CREATE TABLE paimon.delete_inputs.t (id INT) WITH (
+ 'row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true',
+ 'deletion-vectors.enabled' = 'true')""")
+ context.sql('INSERT INTO paimon.delete_inputs.t (id) VALUES (10)')
+ table = PaimonCatalog({'warehouse':
str(tmp_path)}).get_table('delete_inputs.t')
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ update =
builder.new_update().with_plan(table.new_read_builder().new_scan().plan_all())
+ options = {'commit_identifier': 7} if stream else {}
+
+ for empty in [[], pa.array([], type=pa.int8()), pa.chunked_array([],
type=pa.uint64())]:
+ assert update.delete_by_row_id(empty, **options) == []
+ for ids in [
+ (0,),
+ pa.array([0], type=pa.int8()),
+ pa.array([0], type=pa.uint32()),
+ pa.array([None, 0, None], type=pa.int64()).slice(1, 1),
+ pa.chunked_array([pa.array([], type=pa.uint64()), pa.array([0],
type=pa.uint64())]),
+ ]:
+ assert update.delete_by_row_id(ids, **options)
+
+ def index_files():
+ return {path for path in Path(table.location()).rglob('index-*') if
path.is_file()}
+
+ prepared_indexes = index_files()
+ assert prepared_indexes
+ for ids, error in [
+ (pa.array([None], type=pa.int64()), ValueError),
+ (pa.chunked_array([[0], [None]], type=pa.int64()), ValueError),
+ (pa.array([1.5], type=pa.float64()), TypeError),
+ (pa.chunked_array([], type=pa.float64()), TypeError),
+ (pa.array([], type=pa.bool_()), TypeError),
+ (pa.array([], type=pa.string()), TypeError),
+ (pa.array([2 ** 63], type=pa.uint64()), ValueError),
+ (pa.chunked_array([[0], [2 ** 63]], type=pa.uint64()), ValueError),
+ (pa.array([-1], type=pa.int64()), ValueError),
+ ([1.5], TypeError),
+ ([None], TypeError),
+ ([2 ** 63], OverflowError),
+ ]:
+ with pytest.raises(error):
+ update.delete_by_row_id(ids, **options)
+ assert index_files() == prepared_indexes
+ assert table.latest_snapshot().id() == 1
+
+
[email protected]('stream', [False, True])
+def test_sequence_row_id_deletes_do_not_require_pyarrow(tmp_path, stream,
monkeypatch):
+ import builtins
+
+ context = SQLContext()
+ context.register_catalog('paimon', {'warehouse': str(tmp_path)})
+ context.sql('CREATE SCHEMA paimon.sequence_deletes')
+ context.sql("""CREATE TABLE paimon.sequence_deletes.t (id INT) WITH (
+ 'row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true',
+ 'deletion-vectors.enabled' = 'true')""")
+ context.sql('INSERT INTO paimon.sequence_deletes.t (id) VALUES (10), (20),
(30)')
+
+ original_import = builtins.__import__
+
+ def without_pyarrow(name, *args, **kwargs):
+ if name == 'pyarrow' or name.startswith('pyarrow.'):
+ raise ModuleNotFoundError("No module named 'pyarrow'",
name='pyarrow')
+ return original_import(name, *args, **kwargs)
+
+ # PyArrow is optional for the SDK. Plain sequence deletes and commits
+ # must remain usable against existing tables without that dependency.
+ with monkeypatch.context() as patched:
+ patched.setattr(builtins, '__import__', without_pyarrow)
+ table = PaimonCatalog({'warehouse':
str(tmp_path)}).get_table('sequence_deletes.t')
+ builder =
(table.new_stream_write_builder().with_commit_user('sequence-deletes')
+ if stream else table.new_batch_write_builder())
+ for identifier, row_ids in enumerate([[], [0], (1,)]):
+ update = builder.new_update()
+ messages = (update.delete_by_row_id(row_ids, identifier) if stream
+ else update.delete_by_row_id(row_ids))
+ if not row_ids:
+ assert messages == []
+ continue
+ assert messages
+ commit = builder.new_commit()
+ commit.commit(identifier, messages) if stream else
commit.commit(messages)
+
+ assert pa.Table.from_batches(context.sql(
+ 'SELECT id FROM paimon.sequence_deletes.t')).column('id').to_pylist()
== [30]
+
+
def
test_upsert_batches_preserve_missing_fields_nulls_and_stream_updates(tmp_path):
context = SQLContext()
context.register_catalog('paimon', {'warehouse': str(tmp_path)})
diff --git a/crates/paimon/src/table/table_update.rs
b/crates/paimon/src/table/table_update.rs
index d9257a48..a0f69743 100644
--- a/crates/paimon/src/table/table_update.rs
+++ b/crates/paimon/src/table/table_update.rs
@@ -353,8 +353,12 @@ impl TableUpdate {
}
/// Delete rows by row ID using deletion vectors.
- pub async fn delete_by_row_id(&self, row_ids: Vec<i64>) ->
crate::Result<Vec<CommitMessage>> {
- if row_ids.is_empty() {
+ pub async fn delete_by_row_id(
+ &self,
+ row_ids: impl IntoIterator<Item = i64>,
+ ) -> crate::Result<Vec<CommitMessage>> {
+ let mut row_ids = row_ids.into_iter().peekable();
+ if row_ids.peek().is_none() {
return Ok(Vec::new());
}
// The high-level updater selects a read snapshot; the low-level
diff --git a/crates/paimon/src/table/table_update_plan_tests.rs
b/crates/paimon/src/table/table_update_plan_tests.rs
index a86c4fd8..764871cd 100644
--- a/crates/paimon/src/table/table_update_plan_tests.rs
+++ b/crates/paimon/src/table/table_update_plan_tests.rs
@@ -240,7 +240,12 @@ async fn
planned_delete_preserves_unselected_bucket_deletion_vectors() {
update
.with_plan(Plan::from_splits(plan.snapshot_id(), vec![selected]))
.unwrap();
- let messages = update.delete_by_row_id(vec![1]).await.unwrap();
+ append(&table, 30).await;
+ let row_ids = Int64Array::from(vec![-1, 1, 1, -1]).slice(1, 2);
+ let messages = update
+ .delete_by_row_id(row_ids.values().iter().copied())
+ .await
+ .unwrap();
assert_eq!(messages[0].check_from_snapshot, Some(3));
assert_eq!(messages[0].deleted_index_files, vec![previous_index]);
let ranges = messages[0].new_index_files[0]
@@ -265,7 +270,20 @@ async fn
planned_delete_preserves_unselected_bucket_deletion_vectors() {
.try_collect()
.await
.unwrap();
- assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
+ let values: Vec<i32> = batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .iter()
+ .copied()
+ })
+ .collect();
+ assert_eq!(values, vec![30]);
}
#[tokio::test]
diff --git a/crates/paimon/tests/data_evolution_row_sidecar_test.rs
b/crates/paimon/tests/data_evolution_row_sidecar_test.rs
index e25d7fd4..0fc6f89a 100644
--- a/crates/paimon/tests/data_evolution_row_sidecar_test.rs
+++ b/crates/paimon/tests/data_evolution_row_sidecar_test.rs
@@ -561,7 +561,7 @@ async fn
limit_and_deletion_vectors_do_not_trigger_sidecar_reads() {
// Keeping only one live row must not make an otherwise full read use a
// sidecar either. All the sidecars above are absent, but primaries remain.
- let deletions = update.delete_by_row_id((1..100).collect()).await.unwrap();
+ let deletions = update.delete_by_row_id(1..100).await.unwrap();
table
.new_write_builder()
.new_commit()
