| ################################################################################ |
| # 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 os |
| import shutil |
| import tempfile |
| import unittest |
| |
| import pandas |
| import pyarrow as pa |
| |
| from pypaimon import CatalogFactory, Schema |
| |
| from pypaimon.schema.schema_manager import SchemaManager |
| from pypaimon.schema.table_schema import TableSchema |
| |
| |
| class SchemaEvolutionReadTest(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| cls.tempdir = tempfile.mkdtemp() |
| cls.warehouse = os.path.join(cls.tempdir, 'warehouse') |
| cls.catalog = CatalogFactory.create({ |
| 'warehouse': cls.warehouse |
| }) |
| cls.catalog.create_database('default', False) |
| |
| cls.pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('item_id', pa.int64()), |
| ('behavior', pa.string()), |
| ('dt', pa.string()) |
| ]) |
| cls.raw_data = { |
| 'user_id': [1, 2, 3, 4, 5], |
| 'item_id': [1001, 1002, 1003, 1004, 1005], |
| 'behavior': ['a', 'b', 'c', None, 'e'], |
| 'dt': ['p1', 'p1', 'p1', 'p1', 'p2'], |
| } |
| cls.expected = pa.Table.from_pydict(cls.raw_data, schema=cls.pa_schema) |
| |
| @classmethod |
| def tearDownClass(cls): |
| shutil.rmtree(cls.tempdir, ignore_errors=True) |
| |
| def test_schema_evolution(self): |
| # schema 0 |
| pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('item_id', pa.int64()), |
| ('dt', pa.string()) |
| ]) |
| schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.test_sample', schema, False) |
| table1 = self.catalog.get_table('default.test_sample') |
| write_builder = table1.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [1, 2, 3, 4], |
| 'item_id': [1001, 1002, 1003, 1004], |
| 'dt': ['p1', 'p1', 'p2', 'p1'], |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # schema 1 add behavior column |
| pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('item_id', pa.int64()), |
| ('dt', pa.string()), |
| ('behavior', pa.string()) |
| ]) |
| schema2 = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.test_schema_evolution', schema2, False) |
| table2 = self.catalog.get_table('default.test_schema_evolution') |
| table2.table_schema.id = 1 |
| write_builder = table2.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data2 = { |
| 'user_id': [5, 6, 7, 8], |
| 'item_id': [1005, 1006, 1007, 1008], |
| 'dt': ['p2', 'p1', 'p2', 'p2'], |
| 'behavior': ['e', 'f', 'g', 'h'], |
| } |
| pa_table = pa.Table.from_pydict(data2, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # write schema-0 and schema-1 to table2 |
| schema_manager = SchemaManager(table2.file_io, table2.table_path) |
| schema_manager.commit(TableSchema.from_schema(schema_id=0, schema=schema)) |
| schema_manager.commit(TableSchema.from_schema(schema_id=1, schema=schema2)) |
| |
| splits = self._scan_table(table1.new_read_builder()) |
| read_builder = table2.new_read_builder() |
| splits2 = self._scan_table(read_builder) |
| splits.extend(splits2) |
| |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(splits) |
| expected = pa.Table.from_pydict({ |
| 'user_id': [1, 2, 4, 3, 5, 7, 8, 6], |
| 'item_id': [1001, 1002, 1004, 1003, 1005, 1007, 1008, 1006], |
| 'dt': ["p1", "p1", "p1", "p2", "p2", "p2", "p2", "p1"], |
| 'behavior': [None, None, None, None, "e", "g", "h", "f"], |
| }, schema=pa_schema) |
| self.assertEqual(expected, actual) |
| |
| def test_schema_evolution_type(self): |
| # schema 0 |
| pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('time', pa.timestamp('s')), |
| ('dt', pa.string()) |
| ]) |
| schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.schema_evolution_type', schema, False) |
| table1 = self.catalog.get_table('default.schema_evolution_type') |
| write_builder = table1.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [1, 2, 3, 4], |
| 'time': [pandas.Timestamp("2025-01-01 00:00:00"), pandas.Timestamp("2025-01-02 00:02:00"), |
| pandas.Timestamp("2025-01-03 00:03:00"), pandas.Timestamp("2025-01-04 00:04:00")], |
| 'dt': ['p1', 'p1', 'p2', 'p1'], |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # schema 1 add behavior column |
| pa_schema = pa.schema([ |
| ('user_id', pa.int8()), |
| ('time', pa.timestamp('ms')), |
| ('dt', pa.string()), |
| ('behavior', pa.string()) |
| ]) |
| schema2 = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.schema_evolution_type2', schema2, False) |
| table2 = self.catalog.get_table('default.schema_evolution_type2') |
| table2.table_schema.id = 1 |
| write_builder = table2.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data2 = { |
| 'user_id': [5, 6, 7, 8], |
| 'time': [pandas.Timestamp("2025-01-05 00:05:00"), pandas.Timestamp("2025-01-06 00:06:00"), |
| pandas.Timestamp("2025-01-07 00:07:00"), pandas.Timestamp("2025-01-08 00:08:00")], |
| 'dt': ['p2', 'p1', 'p2', 'p2'], |
| 'behavior': ['e', 'f', 'g', 'h'], |
| } |
| pa_table = pa.Table.from_pydict(data2, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # write schema-0 and schema-1 to table2 |
| schema_manager = SchemaManager(table2.file_io, table2.table_path) |
| schema_manager.commit(TableSchema.from_schema(schema_id=0, schema=schema)) |
| schema_manager.commit(TableSchema.from_schema(schema_id=1, schema=schema2)) |
| |
| splits = self._scan_table(table1.new_read_builder()) |
| read_builder = table2.new_read_builder() |
| splits2 = self._scan_table(read_builder) |
| splits.extend(splits2) |
| |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(splits) |
| expected = pa.Table.from_pydict({ |
| 'user_id': [1, 2, 4, 3, 5, 7, 8, 6], |
| 'time': [pandas.Timestamp("2025-01-01 00:00:00"), pandas.Timestamp("2025-01-02 00:02:00"), |
| pandas.Timestamp("2025-01-04 00:04:00"), pandas.Timestamp("2025-01-03 00:03:00"), |
| pandas.Timestamp("2025-01-05 00:05:00"), pandas.Timestamp("2025-01-07 00:07:00"), |
| pandas.Timestamp("2025-01-08 00:08:00"), pandas.Timestamp("2025-01-06 00:06:00"), ], |
| 'dt': ["p1", "p1", "p1", "p2", "p2", "p2", "p2", "p1"], |
| 'behavior': [None, None, None, None, "e", "g", "h", "f"], |
| }, schema=pa_schema) |
| self.assertEqual(expected, actual) |
| |
| def test_schema_evolution_with_scan_filter(self): |
| # schema 0 |
| pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('item_id', pa.int64()), |
| ('dt', pa.string()) |
| ]) |
| schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.test_schema_evolution1', schema, False) |
| table1 = self.catalog.get_table('default.test_schema_evolution1') |
| write_builder = table1.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [1, 2, 3, 4], |
| 'item_id': [1001, 1002, 1003, 1004], |
| 'dt': ['p1', 'p1', 'p2', 'p1'], |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # schema 1 add behavior column |
| pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('item_id', pa.int64()), |
| ('behavior', pa.string()), |
| ('dt', pa.string()) |
| ]) |
| schema2 = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.test_schema_evolution2', schema2, False) |
| table2 = self.catalog.get_table('default.test_schema_evolution2') |
| table2.table_schema.id = 1 |
| write_builder = table2.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data2 = { |
| 'user_id': [5, 6, 7, 8], |
| 'item_id': [1005, 1006, 1007, 1008], |
| 'behavior': ['e', 'f', 'g', 'h'], |
| 'dt': ['p2', 'p1', 'p2', 'p2'], |
| } |
| pa_table = pa.Table.from_pydict(data2, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # write schema-0 and schema-1 to table2 |
| schema_manager = SchemaManager(table2.file_io, table2.table_path) |
| schema_manager.commit(TableSchema.from_schema(schema_id=0, schema=schema)) |
| schema_manager.commit(TableSchema.from_schema(schema_id=1, schema=schema2)) |
| # scan filter for schema evolution |
| latest_snapshot = table1.new_read_builder().new_scan().starting_scanner.snapshot_manager.get_latest_snapshot() |
| table2.table_path = table1.table_path |
| new_read_buidler = table2.new_read_builder() |
| predicate_builder = new_read_buidler.new_predicate_builder() |
| predicate = predicate_builder.less_than('user_id', 3) |
| new_scan = new_read_buidler.with_filter(predicate).new_scan() |
| manifest_files = new_scan.starting_scanner.manifest_list_manager.read_all(latest_snapshot) |
| entries = new_scan.starting_scanner.read_manifest_entries(manifest_files) |
| self.assertEqual(1, len(entries)) # verify scan filter success for schema evolution |
| |
| def test_schema_evolution_with_read_filter(self): |
| # schema 0 |
| pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('item_id', pa.int64()), |
| ('dt', pa.string()) |
| ]) |
| schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.test_schema_evolution_with_filter', schema, False) |
| table1 = self.catalog.get_table('default.test_schema_evolution_with_filter') |
| write_builder = table1.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [1, 2, 3, 4], |
| 'item_id': [1001, 1002, 1003, 1004], |
| 'dt': ['p1', 'p1', 'p2', 'p1'], |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # schema 1 add behavior column |
| pa_schema = pa.schema([ |
| ('user_id', pa.int64()), |
| ('item_id', pa.int64()), |
| ('dt', pa.string()), |
| ('behavior', pa.string()) |
| ]) |
| schema2 = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt']) |
| self.catalog.create_table('default.test_schema_evolution_with_filter2', schema2, False) |
| table2 = self.catalog.get_table('default.test_schema_evolution_with_filter2') |
| table2.table_schema.id = 1 |
| write_builder = table2.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data2 = { |
| 'user_id': [5, 6, 7, 8], |
| 'item_id': [1005, 1006, 1007, 1008], |
| 'dt': ['p2', 'p1', 'p2', 'p2'], |
| 'behavior': ['e', 'f', 'g', 'h'], |
| } |
| pa_table = pa.Table.from_pydict(data2, schema=pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # write schema-0 and schema-1 to table2 |
| schema_manager = SchemaManager(table2.file_io, table2.table_path) |
| schema_manager.commit(TableSchema.from_schema(schema_id=0, schema=schema)) |
| schema_manager.commit(TableSchema.from_schema(schema_id=1, schema=schema2)) |
| |
| # behavior or user_id filter |
| splits = self._scan_table(table1.new_read_builder()) |
| read_builder = table2.new_read_builder() |
| predicate_builder = read_builder.new_predicate_builder() |
| ne_predicate = predicate_builder.equal('behavior', "g") |
| lt_predicate = predicate_builder.less_than('user_id', 6) |
| and_predicate = predicate_builder.or_predicates([ne_predicate, lt_predicate]) |
| splits2 = self._scan_table(read_builder.with_filter(and_predicate)) |
| for split in splits2: |
| for file in split.files: |
| file.schema_id = 1 |
| splits.extend(splits2) |
| |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(splits) |
| expected = pa.Table.from_pydict({ |
| 'user_id': [1, 2, 4, 3, 5, 7], |
| 'item_id': [1001, 1002, 1004, 1003, 1005, 1007], |
| 'dt': ["p1", "p1", "p1", "p2", "p2", "p2"], |
| 'behavior': [None, None, None, None, "e", "g"], |
| }, schema=pa_schema) |
| self.assertEqual(expected, actual) |
| |
| # behavior and user_id filter |
| splits = self._scan_table(table1.new_read_builder()) |
| |
| read_builder = table2.new_read_builder() |
| predicate_builder = read_builder.new_predicate_builder() |
| ne_predicate = predicate_builder.equal('behavior', "g") |
| lt_predicate = predicate_builder.less_than('user_id', 8) |
| and_predicate = predicate_builder.and_predicates([ne_predicate, lt_predicate]) |
| splits2 = self._scan_table(read_builder.with_filter(and_predicate)) |
| for split in splits2: |
| for file in split.files: |
| file.schema_id = 1 |
| splits.extend(splits2) |
| |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(splits) |
| expected = pa.Table.from_pydict({ |
| 'user_id': [1, 2, 4, 3, 7], |
| 'item_id': [1001, 1002, 1004, 1003, 1007], |
| 'dt': ["p1", "p1", "p1", "p2", "p2"], |
| 'behavior': [None, None, None, None, "g"], |
| }, schema=pa_schema) |
| self.assertEqual(expected, actual) |
| |
| # user_id filter |
| splits = self._scan_table(table1.new_read_builder()) |
| |
| read_builder = table2.new_read_builder() |
| predicate_builder = read_builder.new_predicate_builder() |
| predicate = predicate_builder.less_than('user_id', 6) |
| splits2 = self._scan_table(read_builder.with_filter(predicate)) |
| self.assertEqual(1, len(splits2)) |
| for split in splits2: |
| for file in split.files: |
| file.schema_id = 1 |
| splits.extend(splits2) |
| |
| table_read = read_builder.new_read() |
| actual = table_read.to_arrow(splits) |
| expected = pa.Table.from_pydict({ |
| 'user_id': [1, 2, 4, 3, 5], |
| 'item_id': [1001, 1002, 1004, 1003, 1005], |
| 'dt': ["p1", "p1", "p1", "p2", "p2"], |
| 'behavior': [None, None, None, None, "e"], |
| }, schema=pa_schema) |
| self.assertEqual(expected, actual) |
| |
| def _write_test_table(self, table): |
| write_builder = table.new_batch_write_builder() |
| |
| # first write |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [1, 2, 3, 4], |
| 'item_id': [1001, 1002, 1003, 1004], |
| 'behavior': ['a', 'b', 'c', None], |
| 'dt': ['p1', 'p1', 'p2', 'p1'], |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # second write |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data2 = { |
| 'user_id': [5, 6, 7, 8], |
| 'item_id': [1005, 1006, 1007, 1008], |
| 'behavior': ['e', 'f', 'g', 'h'], |
| 'dt': ['p2', 'p1', 'p2', 'p2'], |
| } |
| pa_table = pa.Table.from_pydict(data2, schema=self.pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| def _scan_table(self, read_builder): |
| splits = read_builder.new_scan().plan().splits() |
| return splits |
