| # 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. | |
| import inspect | |
| import os | |
| import shutil | |
| import tempfile | |
| import types | |
| import unittest | |
| import uuid | |
| from unittest import mock | |
| import pyarrow as pa | |
| import pytest | |
| pypaimon = pytest.importorskip("pypaimon") | |
| ray = pytest.importorskip("ray") | |
| from pypaimon import CatalogFactory, Schema | |
| from pypaimon.ray import read_paimon, update_by_row_id | |
| class RayUpdateByRowIdTest(unittest.TestCase): | |
| """Distributed row-id update: rewrite only the files owning the given row ids, | |
| without reading or joining the whole target (unlike merge_into(on=_ROW_ID)).""" | |
| pa_schema = pa.schema([ | |
| ("id", pa.int32()), | |
| ("name", pa.string()), | |
| ("age", pa.int32()), | |
| ]) | |
| de_options = {"row-tracking.enabled": "true", "data-evolution.enabled": "true"} | |
| @classmethod | |
| def setUpClass(cls): | |
| cls.tempdir = tempfile.mkdtemp() | |
| cls.catalog_options = {"warehouse": os.path.join(cls.tempdir, "wh")} | |
| cls.catalog = CatalogFactory.create(cls.catalog_options) | |
| cls.catalog.create_database("default", True) | |
| if not ray.is_initialized(): | |
| ray.init(ignore_reinit_error=True, num_cpus=2) | |
| @classmethod | |
| def tearDownClass(cls): | |
| try: | |
| if ray.is_initialized(): | |
| ray.shutdown() | |
| except Exception: | |
| pass | |
| shutil.rmtree(cls.tempdir, ignore_errors=True) | |
| def _create(self, options=None): | |
| name = f"default.u_{uuid.uuid4().hex[:8]}" | |
| opts = self.de_options if options is None else options | |
| self.catalog.create_table( | |
| name, Schema.from_pyarrow_schema(self.pa_schema, options=opts), False) | |
| return name | |
| def _write(self, target, data): | |
| t = self.catalog.get_table(target) | |
| wb = t.new_batch_write_builder() | |
| w = wb.new_write() | |
| # These fixtures use the file count to derive Ray shuffle parallelism. | |
| chunk_size = min(max(1, data.num_rows), t.options.target_file_row_num()) | |
| for batch in data.to_batches(max_chunksize=chunk_size): | |
| w.write_arrow_batch(batch) | |
| wb.new_commit().commit(w.prepare_commit()) | |
| w.close() | |
| def _read(self, target, projection=None): | |
| t = self.catalog.get_table(target) | |
| rb = t.new_read_builder() | |
| if projection is not None: | |
| rb = rb.with_projection(projection) | |
| return rb.new_read().to_arrow(rb.new_scan().plan().splits()) | |
| def _rowid_by_id(self, target): | |
| tab = self._read(target, ["_ROW_ID", "id"]) | |
| return dict(zip(tab.column("id").to_pylist(), tab.column("_ROW_ID").to_pylist())) | |
| def test_update_by_row_id_basic(self): | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": list(range(1, 7)), "name": [f"n{i}" for i in range(1, 7)], | |
| "age": [i * 10 for i in range(1, 7)]}, schema=self.pa_schema)) | |
| rid = self._rowid_by_id(target) | |
| # update age for ids 2 and 5 only, addressed by their _ROW_ID | |
| src = pa.table({"_ROW_ID": [rid[2], rid[5]], "age": [999, 888]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())])) | |
| # Proof of no full-target read: read_paimon is never called (source is a | |
| # Dataset, and the update routes by manifest metadata, not a scan). | |
| import pypaimon.ray.ray_paimon as rp | |
| with mock.patch.object(rp, "read_paimon", | |
| side_effect=AssertionError("target was read!")): | |
| stats = update_by_row_id(target, ray.data.from_arrow(src), | |
| self.catalog_options, update_cols=["age"]) | |
| self.assertEqual(stats, {"num_updated": 2}) | |
| back = self._read(target).sort_by("id").to_pydict() | |
| self.assertEqual(back["age"], [10, 999, 30, 40, 888, 60]) | |
| self.assertEqual(back["name"], [f"n{i}" for i in range(1, 7)]) # untouched | |
| def test_updates_correct_row_across_files(self): | |
| # A _ROW_ID owned by a middle data file must update only that row. | |
| target = self._create() | |
| for chunk in ([10, 11, 12], [20, 21], [30, 31, 32, 33]): | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": chunk, "name": ["x"] * len(chunk), "age": [0] * len(chunk)}, | |
| schema=self.pa_schema)) | |
| rid = self._rowid_by_id(target) | |
| src = pa.table({"_ROW_ID": [rid[21]], "age": [999]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())])) | |
| stats = update_by_row_id(target, ray.data.from_arrow(src), | |
| self.catalog_options, update_cols=["age"]) | |
| self.assertEqual(stats, {"num_updated": 1}) | |
| back = self._read(target).sort_by("id").to_pydict() | |
| got = dict(zip(back["id"], back["age"])) | |
| self.assertEqual(got[21], 999) | |
| self.assertTrue(all(v == 0 for k, v in got.items() if k != 21)) | |
| def test_empty_dataset_after_transform_is_noop(self): | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) | |
| rid = self._rowid_by_id(target) | |
| source = ray.data.from_arrow(pa.table({ | |
| "_ROW_ID": pa.array([rid[1]], pa.int64()), | |
| "age": pa.array([9], pa.int32()), | |
| })).map_batches( | |
| lambda batch: batch.slice(0, 0), | |
| batch_format="pyarrow", | |
| ) | |
| self.assertEqual( | |
| update_by_row_id( | |
| target, source, self.catalog_options, update_cols=["age"]), | |
| {"num_updated": 0}, | |
| ) | |
| self.assertEqual(self._read(target).column("age").to_pylist(), [1]) | |
| def test_pins_base_snapshot_for_conflict_detection(self): | |
| # The update pins its base snapshot and threads it to distributed_update_apply, | |
| # which uses it for commit-time conflict detection against concurrent writers. | |
| import importlib | |
| m = importlib.import_module("pypaimon.ray.update_by_row_id") # module, not the fn | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]}, schema=self.pa_schema)) | |
| expected_sid = self.catalog.get_table( | |
| target).snapshot_manager().get_latest_snapshot().id | |
| rid = self._rowid_by_id(target) | |
| src = pa.table({"_ROW_ID": [rid[1]], "age": [9]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())])) | |
| captured = {} | |
| def fake_apply(update_ds, table, cols, *, num_partitions, | |
| ray_remote_args=None, base_snapshot_id=None, | |
| estimated_size_bytes=None, estimated_num_rows=None, | |
| data_context=None): | |
| captured["base_snapshot_id"] = base_snapshot_id | |
| captured["num_partitions"] = num_partitions | |
| captured["estimated_size_bytes"] = estimated_size_bytes | |
| captured["estimated_num_rows"] = estimated_num_rows | |
| captured["data_context"] = data_context | |
| return [], 0, [] | |
| with mock.patch.object(m, "distributed_update_apply", fake_apply): | |
| update_by_row_id(target, src, self.catalog_options, update_cols=["age"]) | |
| self.assertEqual(captured["base_snapshot_id"], expected_sid) | |
| self.assertIsNone(captured["num_partitions"]) | |
| self.assertGreater(captured["estimated_size_bytes"], 0) | |
| self.assertEqual(captured["estimated_num_rows"], 1) | |
| self.assertIsNotNone(captured["data_context"]) | |
| def test_transformed_paimon_source_adapts_partitions(self): | |
| import pypaimon.ray.data_evolution_merge_join as merge_join | |
| target = self._create() | |
| for value in range(3): | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [value], "name": ["n{}".format(value)], "age": [0]}, | |
| schema=self.pa_schema, | |
| )) | |
| source = read_paimon( | |
| target, | |
| self.catalog_options, | |
| projection=["_ROW_ID", "age"], | |
| ).map_batches( | |
| lambda batch: pa.table({ | |
| "_ROW_ID": batch.column("_ROW_ID"), | |
| "age": pa.array([99] * batch.num_rows, type=pa.int32()), | |
| }), | |
| batch_format="pyarrow", | |
| ) | |
| real_resolve = merge_join._resolve_row_id_num_partitions | |
| resolved = [] | |
| def track_resolve(*args, **kwargs): | |
| result = real_resolve(*args, **kwargs) | |
| resolved.append((args, result)) | |
| return result | |
| with mock.patch( | |
| "ray.cluster_resources", return_value={"CPU": 320} | |
| ), mock.patch.object( | |
| merge_join, | |
| "_resolve_row_id_num_partitions", | |
| side_effect=track_resolve, | |
| ): | |
| stats = update_by_row_id( | |
| target, | |
| source, | |
| self.catalog_options, | |
| update_cols=["age"], | |
| ) | |
| self.assertEqual(stats, {"num_updated": 3}) | |
| self.assertEqual([result for _, result in resolved], [3]) | |
| cardinality_parameter = inspect.signature( | |
| ray.data.Dataset.map_batches | |
| ).parameters.get("udf_modifying_row_count") | |
| expected_num_rows = ( | |
| 3 | |
| if cardinality_parameter is not None | |
| and cardinality_parameter.default is False | |
| else None | |
| ) | |
| self.assertEqual( | |
| resolved[0][0][1:], (None, expected_num_rows, 3) | |
| ) | |
| self.assertEqual(self._read(target).column("age").to_pylist(), [99] * 3) | |
| def test_expanding_transform_keeps_target_parallelism(self): | |
| import pypaimon.ray.data_evolution_merge_join as merge_join | |
| target = self._create() | |
| for value in range(3): | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [value], "name": ["n{}".format(value)], "age": [0]}, | |
| schema=self.pa_schema, | |
| )) | |
| rows = read_paimon( | |
| target, | |
| self.catalog_options, | |
| projection=["_ROW_ID"], | |
| ).take_all() | |
| row_ids = [row["_ROW_ID"] for row in rows] | |
| map_kwargs = {"batch_format": "pyarrow"} | |
| if "udf_modifying_row_count" in inspect.signature( | |
| ray.data.Dataset.map_batches | |
| ).parameters: | |
| map_kwargs["udf_modifying_row_count"] = True | |
| source = ray.data.from_arrow(pa.table({"seed": [0]})).map_batches( | |
| lambda batch: pa.table({ | |
| "_ROW_ID": pa.array(row_ids, type=pa.int64()), | |
| "age": pa.array([77] * len(row_ids), type=pa.int32()), | |
| }), | |
| **map_kwargs, | |
| ) | |
| real_resolve = merge_join._resolve_row_id_num_partitions | |
| resolved = [] | |
| def track_resolve(*args, **kwargs): | |
| result = real_resolve(*args, **kwargs) | |
| resolved.append((args, result)) | |
| return result | |
| with mock.patch( | |
| "ray.cluster_resources", return_value={"CPU": 320} | |
| ), mock.patch.object( | |
| merge_join, | |
| "_resolve_row_id_num_partitions", | |
| side_effect=track_resolve, | |
| ): | |
| stats = update_by_row_id( | |
| target, | |
| source, | |
| self.catalog_options, | |
| update_cols=["age"], | |
| ) | |
| self.assertEqual(stats, {"num_updated": 3}) | |
| self.assertEqual(resolved, [((None, None, None, 3), 3)]) | |
| self.assertEqual(self._read(target).column("age").to_pylist(), [77] * 3) | |
| def test_filtered_paimon_source_uses_ray_shuffle_default(self): | |
| import pypaimon.ray.data_evolution_merge_join as merge_join | |
| target = "default.u_{}".format(uuid.uuid4().hex[:8]) | |
| schema = pa.schema([ | |
| ("content_key", pa.string()), | |
| ("clip_id", pa.int32()), | |
| ("age", pa.int32()), | |
| ]) | |
| options = dict(self.de_options) | |
| options["target-file-row-num"] = "1" | |
| self.catalog.create_table( | |
| target, | |
| Schema.from_pyarrow_schema(schema, options=options), | |
| False, | |
| ) | |
| matched_rows = 39 | |
| total_rows = 201 | |
| self._write(target, pa.Table.from_pydict({ | |
| "content_key": ["wanted"] * matched_rows | |
| + ["other"] * (total_rows - matched_rows), | |
| "clip_id": [7] * matched_rows | |
| + [8] * (total_rows - matched_rows), | |
| "age": [0] * total_rows, | |
| }, schema=schema)) | |
| table = self.catalog.get_table(target) | |
| predicates = table.new_read_builder().new_predicate_builder() | |
| predicate = predicates.and_predicates([ | |
| predicates.equal("content_key", "wanted"), | |
| predicates.equal("clip_id", 7), | |
| ]) | |
| source = read_paimon( | |
| target, | |
| self.catalog_options, | |
| filter=predicate, | |
| projection=["_ROW_ID", "content_key", "clip_id", "age"], | |
| ).map_batches( | |
| lambda batch: pa.table({ | |
| "_ROW_ID": batch.column("_ROW_ID"), | |
| "age": pa.array([88] * batch.num_rows, type=pa.int32()), | |
| }), | |
| batch_format="pyarrow", | |
| ) | |
| real_resolve = merge_join._resolve_row_id_num_partitions | |
| real_groupby = ray.data.Dataset.groupby | |
| resolved = [] | |
| group_partitions = [] | |
| def track_resolve(*args, **kwargs): | |
| with mock.patch( | |
| "ray.cluster_resources", return_value={"CPU": 320} | |
| ): | |
| result = real_resolve(*args, **kwargs) | |
| resolved.append((args, result)) | |
| return result | |
| def run_small_groupby(dataset, key, num_partitions=None): | |
| # Assert the requested value but keep the test execution small. | |
| group_partitions.append(num_partitions) | |
| return real_groupby(dataset, key, num_partitions=4) | |
| with mock.patch.object( | |
| merge_join, | |
| "_resolve_row_id_num_partitions", | |
| side_effect=track_resolve, | |
| ), mock.patch.object( | |
| ray.data.Dataset, | |
| "groupby", | |
| new=run_small_groupby, | |
| ): | |
| stats = update_by_row_id( | |
| target, | |
| source, | |
| self.catalog_options, | |
| update_cols=["age"], | |
| ) | |
| self.assertEqual(stats, {"num_updated": matched_rows}) | |
| self.assertEqual(resolved, [((None, None, None, total_rows), 200)]) | |
| self.assertEqual(group_partitions, [200]) | |
| ages = self._read(target).column("age").to_pylist() | |
| self.assertEqual(ages.count(88), matched_rows) | |
| self.assertEqual(ages.count(0), total_rows - matched_rows) | |
| def test_new_commit_failure_preserves_pending_messages(self): | |
| err = RuntimeError("new_commit failed") | |
| recorder = {} | |
| with self.assertRaisesRegex(RuntimeError, "new_commit failed"): | |
| self._run_with_fake_commit( | |
| recorder=recorder, | |
| new_commit_errors=[err], | |
| ) | |
| self.assertEqual(recorder["commit_calls"], 0) | |
| self.assertEqual(recorder["abort_calls"], 0) | |
| self.assertEqual(recorder["new_commit_calls"], 1) | |
| def test_commit_failure_does_not_abort_after_commit_started(self): | |
| err = RuntimeError("commit failed") | |
| recorder = {} | |
| with self.assertRaisesRegex(RuntimeError, "commit failed"): | |
| self._run_with_fake_commit( | |
| recorder=recorder, | |
| commit_error=err, | |
| ) | |
| self.assertEqual(recorder["commit_calls"], 1) | |
| self.assertEqual(recorder["abort_calls"], 0) | |
| self.assertEqual(recorder["close_calls"], 1) | |
| def test_close_failure_after_success_warns_and_returns_stats(self): | |
| close_error = RuntimeError("close failed") | |
| with self.assertLogs("pypaimon.ray.row_id_conflict_rewriter", level="WARNING") as logs: | |
| recorder = self._run_with_fake_commit(close_error=close_error) | |
| self.assertEqual(recorder["result"], {"num_updated": 3}) | |
| self.assertEqual(recorder["commit_calls"], 1) | |
| self.assertEqual(recorder["abort_calls"], 0) | |
| self.assertIn( | |
| "Failed to close row-id update commit", | |
| "\n".join(logs.output), | |
| ) | |
| def test_accepts_pyarrow_and_pandas_source(self): | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1, 2], "name": ["a", "b"], "age": [1, 2]}, schema=self.pa_schema)) | |
| rid = self._rowid_by_id(target) | |
| # pyarrow.Table source | |
| update_by_row_id( | |
| target, | |
| pa.table({"_ROW_ID": [rid[1]], "age": [77]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())])), | |
| self.catalog_options, update_cols=["age"]) | |
| self.assertEqual(self._read(target).sort_by("id").to_pydict()["age"], [77, 2]) | |
| # pandas.DataFrame source, updating multiple columns at once | |
| import pandas as pd | |
| update_by_row_id( | |
| target, | |
| pd.DataFrame({"_ROW_ID": pd.array([rid[2]], dtype="int64"), | |
| "name": ["z"], "age": pd.array([88], dtype="int32")}), | |
| self.catalog_options, update_cols=["name", "age"]) | |
| back = self._read(target).sort_by("id").to_pydict() | |
| self.assertEqual(back["age"], [77, 88]) | |
| self.assertEqual(back["name"], ["a", "z"]) | |
| def test_rejects_table_name_source(self): | |
| # A source table's system _ROW_ID is its own, not the target's row ids, so a | |
| # table-name source is rejected rather than silently updating wrong rows. | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(target, "default.some_source", self.catalog_options, | |
| update_cols=["age"]) | |
| def test_rejects_non_data_evolution_table(self): | |
| target = self._create(options={}) # plain append table | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) | |
| src = pa.table({"_ROW_ID": [0], "age": [9]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())])) | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(target, src, self.catalog_options, update_cols=["age"]) | |
| def test_rejects_missing_row_id_column(self): | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) | |
| src = pa.table({"age": [9]}, schema=pa.schema([("age", pa.int32())])) | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(target, src, self.catalog_options, update_cols=["age"]) | |
| def test_rejects_partition_column_update(self): | |
| name = f"default.u_{uuid.uuid4().hex[:8]}" | |
| s = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=["name"], | |
| options=self.de_options) | |
| self.catalog.create_table(name, s, False) | |
| self._write(name, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) | |
| src = pa.table({"_ROW_ID": [0], "name": ["b"]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("name", pa.string())])) | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(name, src, self.catalog_options, update_cols=["name"]) | |
| def test_updates_live_rows_with_deletion_vectors(self): | |
| opts = dict(self.de_options, **{"deletion-vectors.enabled": "true"}) | |
| target = self._create(options=opts) | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1, 2, 3, 4], "name": ["a"] * 4, "age": [10, 20, 30, 40]}, | |
| schema=self.pa_schema)) | |
| table = self.catalog.get_table(target) | |
| wb = table.new_batch_write_builder() | |
| commit = wb.new_commit() | |
| try: | |
| commit.commit(wb.new_update().delete_by_row_id([0, 2])) | |
| finally: | |
| commit.close() | |
| src = pa.table({"_ROW_ID": [1, 3], "age": [99, 88]}) | |
| self.assertEqual(update_by_row_id( | |
| target, src, self.catalog_options, update_cols=["age"]), {"num_updated": 2}) | |
| rb = table.new_read_builder().with_projection(["id", "age", "_ROW_ID"]) | |
| rows = rb.new_read().to_arrow(rb.new_scan().plan().splits()).to_pylist() | |
| self.assertEqual(sorted(rows, key=lambda r: r["id"]), [ | |
| {"id": 2, "age": 99, "_ROW_ID": 1}, {"id": 4, "age": 88, "_ROW_ID": 3}]) | |
| # Deleting after workers have prepared their column files must not | |
| # resurrect the row when the update is committed. | |
| import importlib | |
| module = importlib.import_module("pypaimon.ray.update_by_row_id") | |
| original_commit = module._commit_update_messages | |
| def delete_then_commit(target_table, messages, **kwargs): | |
| builder = target_table.new_batch_write_builder() | |
| concurrent = builder.new_commit() | |
| try: | |
| concurrent.commit(builder.new_update().delete_by_row_id([3])) | |
| finally: | |
| concurrent.close() | |
| original_commit(target_table, messages, **kwargs) | |
| with mock.patch.object(module, "_commit_update_messages", delete_then_commit): | |
| update_by_row_id(target, pa.table({"_ROW_ID": [3], "age": [77]}), | |
| self.catalog_options, update_cols=["age"]) | |
| rb = table.new_read_builder().with_projection(["id", "age", "_ROW_ID"]) | |
| self.assertEqual(rb.new_read().to_arrow(rb.new_scan().plan().splits()).to_pylist(), | |
| [{"id": 2, "age": 99, "_ROW_ID": 1}]) | |
| snapshot = table.snapshot_manager().get_latest_snapshot().id | |
| with self.assertRaisesRegex(ValueError, "deleted _ROW_ID"): | |
| update_by_row_id(target, pa.table({"_ROW_ID": [0, 1], "age": [9, 9]}), | |
| self.catalog_options, update_cols=["age"]) | |
| self.assertEqual(table.snapshot_manager().get_latest_snapshot().id, snapshot) | |
| def test_rejects_blob_column_update(self): | |
| blob_schema = pa.schema([("id", pa.int32()), ("payload", pa.large_binary())]) | |
| name = f"default.u_{uuid.uuid4().hex[:8]}" | |
| self.catalog.create_table( | |
| name, Schema.from_pyarrow_schema(blob_schema, options=self.de_options), False) | |
| self._write(name, pa.Table.from_pydict( | |
| {"id": [1], "payload": pa.array([b"x"], pa.large_binary())}, schema=blob_schema)) | |
| src = pa.table({"_ROW_ID": [0], "payload": pa.array([b"y"], pa.large_binary())}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("payload", pa.large_binary())])) | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(name, src, self.catalog_options, update_cols=["payload"]) | |
| def test_empty_target_foreign_row_id_raises(self): | |
| src = pa.table({"_ROW_ID": [0], "age": [9]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())])) | |
| empty_src = pa.table({"_ROW_ID": pa.array([], pa.int64()), | |
| "age": pa.array([], pa.int32())}) | |
| # (a) never written -> no snapshot | |
| target = self._create() | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(target, src, self.catalog_options, update_cols=["age"]) | |
| # empty source against an empty target is a no-op, not an error | |
| self.assertEqual( | |
| update_by_row_id(target, empty_src, self.catalog_options, update_cols=["age"]), | |
| {"num_updated": 0}) | |
| # (b) written then emptied by overwrite -> snapshot exists but 0 live rows | |
| target2 = self._create() | |
| self._write(target2, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) | |
| wb = self.catalog.get_table(target2).new_batch_write_builder().overwrite() | |
| w = wb.new_write() | |
| w.write_arrow(pa.Table.from_pydict( | |
| {"id": pa.array([], pa.int32()), "name": pa.array([], pa.string()), | |
| "age": pa.array([], pa.int32())}, schema=self.pa_schema)) | |
| wb.new_commit().commit(w.prepare_commit()) | |
| w.close() | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(target2, src, self.catalog_options, update_cols=["age"]) | |
| def test_rejects_unknown_and_empty_update_cols(self): | |
| target = self._create() | |
| self._write(target, pa.Table.from_pydict( | |
| {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) | |
| src = pa.table({"_ROW_ID": [0], "age": [9]}, | |
| schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())])) | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(target, src, self.catalog_options, update_cols=["nope"]) | |
| with self.assertRaises(ValueError): | |
| update_by_row_id(target, src, self.catalog_options, update_cols=[]) | |
| def _run_with_fake_commit(self, *, recorder=None, new_commit_errors=None, | |
| commit_error=None, close_error=None): | |
| import importlib | |
| m = importlib.import_module("pypaimon.ray.update_by_row_id") | |
| if recorder is None: | |
| recorder = {} | |
| recorder.update({ | |
| "msgs": [types.SimpleNamespace(check_from_snapshot=1, new_files=[])], | |
| "new_commit_errors": list(new_commit_errors or []), | |
| "commit_error": commit_error, | |
| "close_error": close_error, | |
| "new_commit_calls": 0, | |
| "commit_calls": 0, | |
| "abort_calls": 0, | |
| "close_calls": 0, | |
| }) | |
| class FakeOptions: | |
| def data_evolution_enabled(self): | |
| return True | |
| def row_tracking_enabled(self): | |
| return True | |
| def deletion_vectors_enabled(self): | |
| return False | |
| def data_evolution_row_id_conflict_rewrite_max_size(self): | |
| return 256 * 1024 * 1024 | |
| def commit_timeout(self): | |
| return 60_000 | |
| class FakeCommit: | |
| def commit(self, msgs): | |
| recorder["commit_calls"] += 1 | |
| recorder["commit_msgs"] = list(msgs) | |
| if recorder["commit_error"] is not None: | |
| raise recorder["commit_error"] | |
| def abort(self, msgs): | |
| recorder["abort_calls"] += 1 | |
| recorder["abort_msgs"] = list(msgs) | |
| def close(self): | |
| recorder["close_calls"] += 1 | |
| if recorder["close_error"] is not None: | |
| raise recorder["close_error"] | |
| class FakeWriteBuilder: | |
| def new_commit(self): | |
| recorder["new_commit_calls"] += 1 | |
| if recorder["new_commit_errors"]: | |
| raise recorder["new_commit_errors"].pop(0) | |
| return FakeCommit() | |
| class FakeTable: | |
| field_names = ["age"] | |
| partition_keys = [] | |
| options = FakeOptions() | |
| table_schema = types.SimpleNamespace(id=0, fields=[ | |
| types.SimpleNamespace( | |
| id=0, | |
| name="age", | |
| type=types.SimpleNamespace(type="INT"), | |
| ) | |
| ]) | |
| @property | |
| def schema_manager(self): | |
| return types.SimpleNamespace(latest=lambda: self.table_schema) | |
| def snapshot_manager(self): | |
| return types.SimpleNamespace(get_latest_snapshot=lambda: types.SimpleNamespace( | |
| id=1, | |
| total_record_count=1, | |
| )) | |
| def new_batch_write_builder(self): | |
| return FakeWriteBuilder() | |
| def copy_without_time_travel(self, options): | |
| return self | |
| class FakeCatalog: | |
| def get_table(self, target): | |
| return FakeTable() | |
| class FakeSource: | |
| def schema(self): | |
| return types.SimpleNamespace(names=["_ROW_ID", "age"]) | |
| def map_batches(self, fn, batch_format=None): | |
| return self | |
| parser_path = ( | |
| "pypaimon.schema.data_types.PyarrowFieldParser.from_paimon_schema" | |
| ) | |
| with mock.patch( | |
| "pypaimon.catalog.catalog_factory.CatalogFactory.create", | |
| return_value=FakeCatalog()), \ | |
| mock.patch.object(m, "_normalize_source", | |
| side_effect=lambda source, catalog_options: source), \ | |
| mock.patch.object(m, "build_update_schema", | |
| return_value=pa.schema([ | |
| ("_ROW_ID", pa.int64()), | |
| ("age", pa.int32()), | |
| ])), \ | |
| mock.patch(parser_path, | |
| return_value=pa.schema([ | |
| ("age", pa.int32()), | |
| ])), \ | |
| mock.patch.object(m, "distributed_update_apply", | |
| return_value=(recorder["msgs"], 3, [])): | |
| recorder["result"] = m.update_by_row_id( | |
| "default.fake", | |
| FakeSource(), | |
| self.catalog_options, | |
| update_cols=["age"], | |
| ) | |
| return recorder | |
| if __name__ == "__main__": | |
| unittest.main() |