| ################################################################################ | |
| # 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 unittest | |
| from unittest.mock import Mock, patch | |
| import pyarrow as pa | |
| from pypaimon.table.data_evolution_merge_into import ( | |
| WhenMatched, | |
| WhenNotMatched, | |
| lit, | |
| source_col, | |
| target_col, | |
| ) | |
| from pypaimon.tests.data_evolution_test_helpers import ( | |
| BatchModeMixin, | |
| DataEvolutionTestBase, | |
| StreamModeMixin, | |
| ) | |
| try: | |
| import datafusion # noqa: F401 | |
| _HAS_DATAFUSION = True | |
| except ImportError: | |
| _HAS_DATAFUSION = False | |
| _SKIP_CONDITION = not _HAS_DATAFUSION | |
| _SKIP_REASON = "pypaimon[datafusion] is required for condition expressions" | |
| class TableMergeIntoTest(BatchModeMixin, DataEvolutionTestBase, unittest.TestCase): | |
| def test_missing_datafusion_error_suggests_datafusion_extra(self): | |
| from pypaimon.ray.merge_condition import _load_datafusion | |
| with patch.dict("sys.modules", {"datafusion": None}): | |
| with self.assertRaisesRegex( | |
| ImportError, r"pip install 'pypaimon\[datafusion\]'"): | |
| _load_datafusion() | |
| def test_datafusion_conditions_require_python_310(self): | |
| from pypaimon.ray.merge_condition import _load_datafusion | |
| with patch("sys.version_info", (3, 9)): | |
| with self.assertRaisesRegex(ImportError, "Python 3.10 or newer"): | |
| _load_datafusion() | |
| def _read_sorted(self, table): | |
| return self._read_all(table).sort_by("id").to_pydict() | |
| def _read_projected_sorted(self, table, projection): | |
| rb = table.new_read_builder().with_projection(projection) | |
| return rb.new_read().to_arrow(rb.new_scan().plan().splits()).sort_by( | |
| "id" | |
| ).to_pydict() | |
| def _row_ids_by_id(self, table): | |
| rows = self._read_projected_sorted(table, ["id", "_ROW_ID"]) | |
| return dict(zip(rows["id"], rows["_ROW_ID"])) | |
| def _merge_and_commit(self, table, source, **kwargs): | |
| wb = table.new_batch_write_builder() | |
| table_update = wb.new_update() | |
| msgs = table_update.merge_into(source, **kwargs) | |
| commit = wb.new_commit() | |
| commit.commit(msgs) | |
| commit.close() | |
| return msgs | |
| def test_when_matched_action_constructors(self): | |
| update = WhenMatched.update( | |
| {"age": source_col("age")}, | |
| condition="s.age > t.age", | |
| ) | |
| delete = WhenMatched.delete(condition="s.deleted = TRUE") | |
| self.assertEqual({"age": source_col("age")}, update.update) | |
| self.assertEqual("s.age > t.age", update.condition) | |
| self.assertFalse(update.delete) | |
| self.assertIsNone(delete.update) | |
| self.assertEqual("s.deleted = TRUE", delete.condition) | |
| self.assertTrue(delete.delete) | |
| def test_ray_merge_rejects_mixed_update_delete_duplicate_row_ids(self): | |
| import pypaimon.ray.data_evolution_merge_into as ray_merge | |
| with self.assertRaisesRegex(ValueError, "multiple source rows"): | |
| ray_merge._validate_disjoint_action_row_ids([7], [7]) | |
| def _mock_ray_commit_table(self): | |
| table = Mock() | |
| write_builder = Mock() | |
| table_commit = Mock() | |
| table.new_batch_write_builder.return_value = write_builder | |
| write_builder.new_commit.return_value = table_commit | |
| return table, table_commit | |
| def test_ray_execute_validates_mixed_update_delete_duplicate_row_ids(self): | |
| import pypaimon.ray.data_evolution_merge_into as ray_merge | |
| update_msg = object() | |
| delete_msg = object() | |
| table, _ = self._mock_ray_commit_table() | |
| with patch.object( | |
| ray_merge, | |
| "distributed_update_apply", | |
| return_value=([update_msg], 1, [7]) | |
| ), patch.object( | |
| ray_merge, | |
| "distributed_delete_apply", | |
| return_value=([delete_msg], 1, [7]) | |
| ): | |
| with self.assertRaisesRegex(ValueError, "multiple source rows"): | |
| ray_merge._execute_and_commit( | |
| table=table, | |
| update_ds=object(), | |
| delete_ds=object(), | |
| insert_ds=None, | |
| update_cols_union=["age"], | |
| base_snapshot=None, | |
| num_partitions=1, | |
| ray_remote_args=None, | |
| concurrency=None, | |
| ) | |
| table.new_batch_write_builder.assert_not_called() | |
| def test_ray_execute_preserves_messages_on_later_branch_failure(self): | |
| import pypaimon.ray.data_evolution_merge_into as ray_merge | |
| update_msg = object() | |
| table, _ = self._mock_ray_commit_table() | |
| with patch.object( | |
| ray_merge, | |
| "distributed_update_apply", | |
| return_value=([update_msg], 1, []) | |
| ), patch.object( | |
| ray_merge, | |
| "distributed_delete_apply", | |
| side_effect=RuntimeError("delete failed") | |
| ): | |
| with self.assertRaisesRegex(RuntimeError, "delete failed"): | |
| ray_merge._execute_and_commit( | |
| table=table, | |
| update_ds=object(), | |
| delete_ds=object(), | |
| insert_ds=None, | |
| update_cols_union=["age"], | |
| base_snapshot=None, | |
| num_partitions=1, | |
| ray_remote_args=None, | |
| concurrency=None, | |
| ) | |
| table.new_batch_write_builder.assert_not_called() | |
| def test_ray_execute_preserves_messages_on_insert_failure(self): | |
| import pypaimon.ray.data_evolution_merge_into as ray_merge | |
| update_msg = object() | |
| table, _ = self._mock_ray_commit_table() | |
| with patch.object( | |
| ray_merge, | |
| "distributed_update_apply", | |
| return_value=([update_msg], 1, []) | |
| ), patch.object( | |
| ray_merge, | |
| "distributed_write_collect_msgs", | |
| side_effect=RuntimeError("insert failed") | |
| ): | |
| with self.assertRaisesRegex(RuntimeError, "insert failed"): | |
| ray_merge._execute_and_commit( | |
| table=table, | |
| update_ds=object(), | |
| delete_ds=None, | |
| insert_ds=object(), | |
| update_cols_union=["age"], | |
| base_snapshot=None, | |
| num_partitions=1, | |
| ray_remote_args=None, | |
| concurrency=None, | |
| ) | |
| table.new_batch_write_builder.assert_not_called() | |
| def test_ray_execute_does_not_abort_after_commit_starts(self): | |
| import pypaimon.ray.data_evolution_merge_into as ray_merge | |
| update_msg = object() | |
| table, table_commit = self._mock_ray_commit_table() | |
| table_commit.commit.side_effect = RuntimeError("commit failed") | |
| with patch.object( | |
| ray_merge, | |
| "distributed_update_apply", | |
| return_value=([update_msg], 1, []) | |
| ): | |
| with self.assertRaisesRegex(RuntimeError, "commit failed"): | |
| ray_merge._execute_and_commit( | |
| table=table, | |
| update_ds=object(), | |
| delete_ds=None, | |
| insert_ds=None, | |
| update_cols_union=["age"], | |
| base_snapshot=None, | |
| num_partitions=1, | |
| ray_remote_args=None, | |
| concurrency=None, | |
| ) | |
| table_commit.commit.assert_called_once_with([update_msg]) | |
| table_commit.abort.assert_not_called() | |
| table_commit.close.assert_called_once_with() | |
| def test_table_merge_into_updates_and_inserts(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["Alice", "Bob"], | |
| "age": pa.array([25, 30], type=pa.int32()), | |
| "city": ["NYC", "LA"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([2, 3], type=pa.int32()), | |
| "name": ["Bobby", "Cindy"], | |
| "age": pa.array([31, 22], type=pa.int32()), | |
| "city": ["LA2", "SF"], | |
| }, schema=self.pa_schema) | |
| msgs = self._merge_and_commit( | |
| target, | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2, 3], | |
| "name": ["Alice", "Bobby", "Cindy"], | |
| "age": [25, 31, 22], | |
| "city": ["NYC", "LA2", "SF"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| def test_table_merge_into_insert_only_on_empty_table(self): | |
| target = self._create_table() | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["Alice", "Bob"], | |
| "age": pa.array([25, 30], type=pa.int32()), | |
| "city": ["NYC", "LA"], | |
| }, schema=self.pa_schema) | |
| msgs = self._merge_and_commit( | |
| target, | |
| source, | |
| on=["id"], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2], | |
| "name": ["Alice", "Bob"], | |
| "age": [25, 30], | |
| "city": ["NYC", "LA"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| def test_table_merge_into_with_renamed_on_key_fills_insert_key(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "source_id": pa.array([1, 2], type=pa.int32()), | |
| "new_name": ["Alicia", "Bob"], | |
| "new_age": pa.array([26, 30], type=pa.int32()), | |
| "new_city": ["SF", "LA"], | |
| }) | |
| self._merge_and_commit( | |
| target, | |
| source, | |
| on={"id": "source_id"}, | |
| when_matched=[ | |
| WhenMatched.update({ | |
| "name": source_col("new_name"), | |
| "age": source_col("new_age"), | |
| "city": source_col("new_city"), | |
| }), | |
| ], | |
| when_not_matched=[ | |
| WhenNotMatched(insert={ | |
| "name": source_col("new_name"), | |
| "age": source_col("new_age"), | |
| "city": source_col("new_city"), | |
| }), | |
| ], | |
| ) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2], | |
| "name": ["Alicia", "Bob"], | |
| "age": [26, 30], | |
| "city": ["SF", "LA"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| def test_table_merge_into_ignores_target_self_assignment(self): | |
| for assignment in [target_col("name"), "t.name"]: | |
| with self.subTest(assignment=assignment): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| row_ids_before = self._row_ids_by_id(target) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alicia"], | |
| "age": pa.array([26], type=pa.int32()), | |
| "city": ["SF"], | |
| }, schema=self.pa_schema) | |
| msgs = self._merge_and_commit( | |
| target, | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update({"name": assignment})], | |
| ) | |
| self.assertEqual([], msgs) | |
| self.assertEqual( | |
| { | |
| "id": [1], | |
| "name": ["Alice"], | |
| "age": [25], | |
| "city": ["NYC"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| self.assertEqual(row_ids_before, self._row_ids_by_id(target)) | |
| def test_table_merge_into_source_assignment_same_name_is_modified(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alicia"], | |
| "age": pa.array([26], type=pa.int32()), | |
| "city": ["SF"], | |
| }, schema=self.pa_schema) | |
| msgs = self._merge_and_commit( | |
| target, | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update({"name": source_col("name")})], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual(["Alicia"], self._read_sorted(target)["name"]) | |
| self.assertEqual([25], self._read_sorted(target)["age"]) | |
| def test_table_merge_into_rejects_nested_assignment_key(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alicia"], | |
| "age": pa.array([26], type=pa.int32()), | |
| "city": ["SF"], | |
| }, schema=self.pa_schema) | |
| with self.assertRaisesRegex(ValueError, "unknown target column"): | |
| target.new_batch_write_builder().new_update().merge_into( | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update({"name.first": lit("A")})], | |
| ) | |
| def test_table_merge_into_rejects_partition_column_update(self): | |
| target = self._create_table(partition_keys=["city"]) | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([26], type=pa.int32()), | |
| "city": ["SF"], | |
| }, schema=self.pa_schema) | |
| with self.assertRaisesRegex(ValueError, "partition columns"): | |
| target.new_batch_write_builder().new_update().merge_into( | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update({"city": source_col("city")})], | |
| ) | |
| def test_table_merge_into_preserves_row_ids_for_matched_rows(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["Alice", "Bob"], | |
| "age": pa.array([25, 30], type=pa.int32()), | |
| "city": ["NYC", "LA"], | |
| }, schema=self.pa_schema)) | |
| row_ids_before = self._row_ids_by_id(target) | |
| self._merge_and_commit( | |
| target, | |
| pa.Table.from_pydict({ | |
| "id": pa.array([2, 3], type=pa.int32()), | |
| "name": ["Bobby", "Cindy"], | |
| "age": pa.array([31, 22], type=pa.int32()), | |
| "city": ["LA2", "SF"], | |
| }, schema=self.pa_schema), | |
| on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| ) | |
| row_ids_after = self._row_ids_by_id(target) | |
| self.assertEqual(row_ids_before[1], row_ids_after[1]) | |
| self.assertEqual(row_ids_before[2], row_ids_after[2]) | |
| self.assertGreater(row_ids_after[3], max(row_ids_before.values())) | |
| def test_table_merge_into_delete_matched_rows(self): | |
| options = dict(self.table_options) | |
| options["deletion-vectors.enabled"] = "true" | |
| target = self._create_table(options=options) | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1, 2, 3], type=pa.int32()), | |
| "name": ["Alice", "Bob", "Charlie"], | |
| "age": pa.array([25, 30, 35], type=pa.int32()), | |
| "city": ["NYC", "LA", "Chicago"], | |
| }, schema=self.pa_schema)) | |
| msgs = self._merge_and_commit( | |
| target, | |
| pa.Table.from_pydict({ | |
| "id": pa.array([2, 3], type=pa.int32()), | |
| "name": ["ignored", "ignored"], | |
| "age": pa.array([99, 99], type=pa.int32()), | |
| "city": ["ignored", "ignored"], | |
| }, schema=self.pa_schema), | |
| on=["id"], | |
| when_matched=[WhenMatched.delete()], | |
| ) | |
| self.assertEqual(1, sum(len(m.index_adds) for m in msgs)) | |
| self.assertEqual( | |
| { | |
| "id": [1], | |
| "name": ["Alice"], | |
| "age": [25], | |
| "city": ["NYC"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| def test_table_merge_into_second_round_updates_first_round_insert(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| self._merge_and_commit( | |
| target, | |
| pa.Table.from_pydict({ | |
| "id": pa.array([2], type=pa.int32()), | |
| "name": ["Bob"], | |
| "age": pa.array([30], type=pa.int32()), | |
| "city": ["LA"], | |
| }, schema=self.pa_schema), | |
| on=["id"], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| ) | |
| row_ids_after_insert = self._row_ids_by_id(target) | |
| self._merge_and_commit( | |
| target, | |
| pa.Table.from_pydict({ | |
| "id": pa.array([2], type=pa.int32()), | |
| "name": ["Bobby"], | |
| "age": pa.array([31], type=pa.int32()), | |
| "city": ["LA2"], | |
| }, schema=self.pa_schema), | |
| on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| ) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2], | |
| "name": ["Alice", "Bobby"], | |
| "age": [25, 31], | |
| "city": ["NYC", "LA2"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| self.assertEqual(row_ids_after_insert, self._row_ids_by_id(target)) | |
| def test_table_merge_into_accepts_paimon_source_table(self): | |
| target = self._create_table() | |
| source = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| self._write_arrow(source, pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["Alicia", "Bob"], | |
| "age": pa.array([26, 30], type=pa.int32()), | |
| "city": ["SF", "LA"], | |
| }, schema=self.pa_schema)) | |
| msgs = self._merge_and_commit( | |
| target, | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2], | |
| "name": ["Alicia", "Bob"], | |
| "age": [26, 30], | |
| "city": ["SF", "LA"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| def test_table_merge_into_updates_and_inserts_blob_column(self): | |
| blob_schema = pa.schema([ | |
| ("id", pa.int32()), | |
| ("name", pa.string()), | |
| ("payload", pa.large_binary()), | |
| ]) | |
| target = self._create_table(pa_schema=blob_schema) | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["Alice", "Bob"], | |
| "payload": pa.array( | |
| [b"blob-1", b"blob-2"], type=pa.large_binary() | |
| ), | |
| }, schema=blob_schema)) | |
| msgs = self._merge_and_commit( | |
| target, | |
| pa.Table.from_pydict({ | |
| "id": pa.array([2, 3], type=pa.int32()), | |
| "name": ["Bobby", "Cindy"], | |
| "payload": pa.array( | |
| [b"blob-2-updated", b"blob-3"], type=pa.large_binary() | |
| ), | |
| }, schema=blob_schema), | |
| on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2, 3], | |
| "name": ["Alice", "Bobby", "Cindy"], | |
| "payload": [b"blob-1", b"blob-2-updated", b"blob-3"], | |
| }, | |
| self._read_projected_sorted(target, ["id", "name", "payload"]), | |
| ) | |
| def test_table_merge_into_nested_payloads(self): | |
| cases = [ | |
| (pa.list_(pa.float32(), 2), [1., 2.], [3., 4.], [5., 6.]), | |
| (pa.list_(pa.int64()), [1, None], [], None), | |
| (pa.map_(pa.string(), pa.string()), [("old", "value")], [], [("new", None)]), | |
| (pa.struct([("label", pa.string()), ("count", pa.int64())]), | |
| {"label": "old", "count": None}, None, {"label": None, "count": 2}), | |
| ] | |
| for payload_type, before, updated, inserted in cases: | |
| with self.subTest(payload_type=payload_type): | |
| schema = pa.schema([("id", pa.int32()), ("payload", payload_type)]) | |
| target = self._create_table(pa_schema=schema, options=dict( | |
| self.table_options, **{ | |
| "file.format": "parquet", "vector.file.format": "parquet", | |
| "deletion-vectors.enabled": "true", | |
| })) | |
| self._write_arrow(target, pa.Table.from_pylist([ | |
| {"id": 1, "payload": before}, {"id": 2, "payload": before}, | |
| ], schema=schema)) | |
| # Keep the source chunked, with a nonzero slice offset. | |
| source = pa.concat_tables([ | |
| pa.Table.from_pylist([ | |
| {"id": 0, "payload": before}, {"id": 2, "payload": updated}, | |
| ], schema=schema).slice(1), | |
| pa.Table.from_pylist([{"id": 3, "payload": inserted}], schema=schema), | |
| ]) | |
| self._merge_and_commit( | |
| target, source, on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| when_not_matched=[WhenNotMatched(insert="*")]) | |
| self.assertEqual({"id": [1, 2, 3], "payload": [before, updated, inserted]}, | |
| self._read_sorted(target)) | |
| # A delete must also accept unused nested source columns. | |
| self._merge_and_commit( | |
| target, source.slice(0, 1), on=["id"], | |
| when_matched=[WhenMatched.delete()]) | |
| self.assertEqual({"id": [1, 3], "payload": [before, inserted]}, | |
| self._read_sorted(target)) | |
| @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON) | |
| def test_nested_merge_keeps_mapped_composite_keys_and_null_semantics(self): | |
| schema = pa.schema([ | |
| ("id", pa.int32()), ("group", pa.string()), | |
| ("payload", pa.list_(pa.int32())), ("previous", pa.list_(pa.int32())), | |
| ]) | |
| target = self._create_table(pa_schema=schema, options=dict( | |
| self.table_options, **{"file.format": "parquet"})) | |
| self._write_arrow(target, pa.Table.from_pylist([ | |
| {"id": 1, "group": "a", "payload": [10], "previous": []}, | |
| {"id": 1, "group": "b", "payload": [20], "previous": []}, | |
| {"id": None, "group": "a", "payload": [30], "previous": []}, | |
| ], schema=schema)) | |
| source_schema = pa.schema([ | |
| ("source_id", pa.int32()), ("source_group", pa.string()), | |
| ("payload", pa.list_(pa.int32())), | |
| ]) | |
| source = pa.Table.from_pylist([ | |
| {"source_id": 1, "source_group": "a", "payload": [11]}, | |
| {"source_id": 1, "source_group": "b", "payload": [21]}, | |
| {"source_id": None, "source_group": "a", "payload": [31]}, | |
| {"source_id": 2, "source_group": "a", "payload": [40]}, | |
| ], schema=source_schema) | |
| self._merge_and_commit( | |
| target, source, on={"id": "source_id", "group": "source_group"}, | |
| when_matched=[WhenMatched.update({ | |
| "payload": source_col("payload"), "previous": target_col("payload"), | |
| }, condition="s.source_id = 1 AND t.group = 'a'")], | |
| when_not_matched=[WhenNotMatched(insert={"payload": source_col("payload")})]) | |
| self.assertCountEqual([ | |
| {"id": 1, "group": "a", "payload": [11], "previous": [10]}, | |
| {"id": 1, "group": "b", "payload": [20], "previous": []}, | |
| {"id": None, "group": "a", "payload": [30], "previous": []}, | |
| {"id": None, "group": "a", "payload": [31], "previous": None}, | |
| {"id": 2, "group": "a", "payload": [40], "previous": None}, | |
| ], self._read_all(target).to_pylist()) | |
| def test_nested_merge_rejects_duplicate_matches_before_commit(self): | |
| schema = pa.schema([("id", pa.int32()), ("payload", pa.list_(pa.int32()))]) | |
| target = self._create_table(pa_schema=schema, options=dict( | |
| self.table_options, **{"file.format": "parquet"})) | |
| self._write_arrow(target, pa.Table.from_pylist([ | |
| {"id": 1, "payload": [10]}, | |
| ], schema=schema)) | |
| snapshot_id = target.snapshot_manager().get_latest_snapshot().id | |
| source = pa.Table.from_pylist([ | |
| {"id": 1, "payload": [11]}, {"id": 1, "payload": [12]}, | |
| {"id": 2, "payload": [20]}, | |
| ], schema=schema) | |
| with self.assertRaisesRegex(ValueError, "multiple source rows"): | |
| self._merge_and_commit( | |
| target, source, on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| when_not_matched=[WhenNotMatched(insert="*")]) | |
| self.assertEqual(snapshot_id, target.snapshot_manager().get_latest_snapshot().id) | |
| self.assertEqual({"id": [1], "payload": [[10]]}, self._read_sorted(target)) | |
| def test_table_merge_into_inserts_null_for_unspecified_blob_column(self): | |
| blob_schema = pa.schema([ | |
| ("id", pa.int32()), | |
| ("name", pa.string()), | |
| ("payload", pa.large_binary()), | |
| ]) | |
| source_schema = pa.schema([ | |
| ("id", pa.int32()), | |
| ("name", pa.string()), | |
| ]) | |
| target = self._create_table(pa_schema=blob_schema) | |
| msgs = self._merge_and_commit( | |
| target, | |
| pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| }, schema=source_schema), | |
| on=["id"], | |
| when_not_matched=[WhenNotMatched(insert={ | |
| "id": source_col("id"), | |
| "name": source_col("name"), | |
| })], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual( | |
| { | |
| "id": [1], | |
| "name": ["Alice"], | |
| "payload": [None], | |
| }, | |
| self._read_projected_sorted(target, ["id", "name", "payload"]), | |
| ) | |
| @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON) | |
| def test_table_merge_into_condition_clauses(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["old_1", "old_2"], | |
| "age": pa.array([10, 40], type=pa.int32()), | |
| "city": ["NYC", "LA"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1, 2, 3], type=pa.int32()), | |
| "name": ["new_1", "new_2", "new_3"], | |
| "age": pa.array([15, 35, 50], type=pa.int32()), | |
| "city": ["SF", "SEA", "DAL"], | |
| }, schema=self.pa_schema) | |
| msgs = self._merge_and_commit( | |
| target, | |
| source, | |
| on=["id"], | |
| when_matched=[ | |
| WhenMatched.update({"age": lit(99)}, condition="s.age > t.age"), | |
| WhenMatched.update({"name": lit("kept")}), | |
| ], | |
| when_not_matched=[ | |
| WhenNotMatched(insert="*", condition="s.age > 45"), | |
| ], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2, 3], | |
| "name": ["old_1", "kept", "new_3"], | |
| "age": [99, 40, 50], | |
| "city": ["NYC", "LA", "DAL"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| def test_table_merge_into_duplicate_source_match_rejected(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1, 1], type=pa.int32()), | |
| "name": ["A1", "A2"], | |
| "age": pa.array([26, 27], type=pa.int32()), | |
| "city": ["SF", "LA"], | |
| }, schema=self.pa_schema) | |
| with self.assertRaisesRegex(ValueError, "multiple source rows"): | |
| target.new_batch_write_builder().new_update().merge_into( | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update({"age": "s.age"})], | |
| ) | |
| def test_table_merge_into_self_merge_by_row_id(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["Alice", "Bob"], | |
| "age": pa.array([25, 30], type=pa.int32()), | |
| "city": ["NYC", "LA"], | |
| }, schema=self.pa_schema)) | |
| msgs = self._merge_and_commit( | |
| target, | |
| target, | |
| on=["_ROW_ID"], | |
| when_matched=[WhenMatched.update({"name": lit("updated")})], | |
| ) | |
| self.assertTrue(msgs) | |
| self.assertEqual(["updated", "updated"], self._read_sorted(target)["name"]) | |
| class StreamTableMergeIntoTest(StreamModeMixin, DataEvolutionTestBase, unittest.TestCase): | |
| def _read_sorted(self, table): | |
| return self._read_all(table).sort_by("id").to_pydict() | |
| def test_stream_table_update_merge_into(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| source = pa.Table.from_pydict({ | |
| "id": pa.array([1, 2], type=pa.int32()), | |
| "name": ["Alicia", "Bob"], | |
| "age": pa.array([26, 30], type=pa.int32()), | |
| "city": ["SF", "LA"], | |
| }, schema=self.pa_schema) | |
| base_snapshot_id = self._latest_snapshot_id(target) | |
| wb = target.new_stream_write_builder() | |
| cid = self._next_commit_id() | |
| msgs = wb.new_update().merge_into( | |
| source, | |
| on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| commit_identifier=cid, | |
| ) | |
| commit = wb.new_commit() | |
| commit.commit(msgs, cid) | |
| commit.close() | |
| self._assert_stream_builder_snapshots( | |
| target, wb, base_snapshot_id, [cid] | |
| ) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2], | |
| "name": ["Alicia", "Bob"], | |
| "age": [26, 30], | |
| "city": ["SF", "LA"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| def test_stream_table_update_merge_into_multiple_rounds(self): | |
| target = self._create_table() | |
| self._write_arrow(target, pa.Table.from_pydict({ | |
| "id": pa.array([1], type=pa.int32()), | |
| "name": ["Alice"], | |
| "age": pa.array([25], type=pa.int32()), | |
| "city": ["NYC"], | |
| }, schema=self.pa_schema)) | |
| base_snapshot_id = self._latest_snapshot_id(target) | |
| wb = target.new_stream_write_builder() | |
| commit = wb.new_commit() | |
| cid1 = self._next_commit_id() | |
| msgs1 = wb.new_update().merge_into( | |
| pa.Table.from_pydict({ | |
| "id": pa.array([2], type=pa.int32()), | |
| "name": ["Bob"], | |
| "age": pa.array([30], type=pa.int32()), | |
| "city": ["LA"], | |
| }, schema=self.pa_schema), | |
| on=["id"], | |
| when_not_matched=[WhenNotMatched(insert="*")], | |
| commit_identifier=cid1, | |
| ) | |
| commit.commit(msgs1, cid1) | |
| cid2 = self._next_commit_id() | |
| msgs2 = wb.new_update().merge_into( | |
| pa.Table.from_pydict({ | |
| "id": pa.array([2], type=pa.int32()), | |
| "name": ["Bobby"], | |
| "age": pa.array([31], type=pa.int32()), | |
| "city": ["LA2"], | |
| }, schema=self.pa_schema), | |
| on=["id"], | |
| when_matched=[WhenMatched.update("*")], | |
| commit_identifier=cid2, | |
| ) | |
| commit.commit(msgs2, cid2) | |
| commit.close() | |
| self._assert_stream_builder_snapshots( | |
| target, wb, base_snapshot_id, [cid1, cid2] | |
| ) | |
| self.assertEqual( | |
| { | |
| "id": [1, 2], | |
| "name": ["Alice", "Bobby"], | |
| "age": [25, 31], | |
| "city": ["NYC", "LA2"], | |
| }, | |
| self._read_sorted(target), | |
| ) | |
| if __name__ == "__main__": | |
| unittest.main() |