| ################################################################################ |
| # 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[sql] is required for condition expressions" |
| |
| |
| class TableMergeIntoTest(BatchModeMixin, DataEvolutionTestBase, unittest.TestCase): |
| |
| 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, table_commit = 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_commit.abort.assert_called_once_with([update_msg, delete_msg]) |
| table_commit.close.assert_called_once_with() |
| |
| def test_ray_execute_aborts_prepared_messages_on_later_branch_failure(self): |
| import pypaimon.ray.data_evolution_merge_into as ray_merge |
| |
| update_msg = object() |
| table, table_commit = 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_commit.abort.assert_called_once_with([update_msg]) |
| table_commit.close.assert_called_once_with() |
| table_commit.commit.assert_not_called() |
| |
| def test_ray_execute_aborts_prepared_messages_on_insert_failure(self): |
| import pypaimon.ray.data_evolution_merge_into as ray_merge |
| |
| update_msg = object() |
| table, table_commit = 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_commit.abort.assert_called_once_with([update_msg]) |
| table_commit.close.assert_called_once_with() |
| table_commit.commit.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_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() |
