| # 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 |
| import pytest |
| |
| from pypaimon import CatalogFactory, Schema |
| |
| from pypaimon.schema.data_types import AtomicType |
| from pypaimon.schema.schema_change import Move, SchemaChange |
| 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) |
| |
| @pytest.mark.python_plan |
| 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) |
| |
| @pytest.mark.python_plan |
| 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_type_promotion_unpartitioned(self): |
| # End-to-end via public API only (create -> write -> alter column type |
| # -> write -> read). A non-partitioned table whose read needs no column |
| # reordering takes the reader fast path that skips partition padding and |
| # index remapping. The file written before the type change keeps its old |
| # physical type, so it must be aligned to the promoted type to |
| # concatenate with the file written after; otherwise the read crashes |
| # with an Arrow schema mismatch. This is not specific to INT -> BIGINT: |
| # it applies to any type change of an existing column -- integer/float |
| # widening, DECIMAL precision/scale changes, and cross-type changes. |
| import decimal |
| |
| # Each case: (name, old arrow type, new arrow type, new Paimon type, |
| # value written to the old-schema file, that same value as it should |
| # read back under the new type, value written to the new-schema file). |
| # The old write value and its expected read form differ for cross-type |
| # changes, where the old file is materialized under the new type. |
| cases = [ |
| ("smallint_to_int", pa.int16(), pa.int32(), 'INT', |
| [10, 20], [10, 20], [30, 40]), |
| ("int_to_bigint", pa.int32(), pa.int64(), 'BIGINT', |
| [10, 20], [10, 20], [30, 40]), |
| ("float_to_double", pa.float32(), pa.float64(), 'DOUBLE', |
| [1.5, 2.5], [1.5, 2.5], [3.5, 4.5]), |
| ("decimal_precision_up", |
| pa.decimal128(10, 2), pa.decimal128(20, 2), 'DECIMAL(20, 2)', |
| [decimal.Decimal('1.23'), decimal.Decimal('4.56')], |
| [decimal.Decimal('1.23'), decimal.Decimal('4.56')], |
| [decimal.Decimal('7.89'), decimal.Decimal('0.12')]), |
| ("decimal_scale_up", |
| pa.decimal128(10, 2), pa.decimal128(10, 4), 'DECIMAL(10, 4)', |
| [decimal.Decimal('1.23'), decimal.Decimal('4.56')], |
| [decimal.Decimal('1.2300'), decimal.Decimal('4.5600')], |
| [decimal.Decimal('7.8901'), decimal.Decimal('0.1234')]), |
| ("int_to_string", pa.int32(), pa.string(), 'STRING', |
| [10, 20], ['10', '20'], ['a', 'b']), |
| # Lossy cross-type change: DOUBLE -> INT truncates (matches Java |
| # CastExecutors), so 1.2/2.8 read back as 1/2. |
| ("double_to_int", pa.float64(), pa.int32(), 'INT', |
| [1.2, 2.8], [1, 2], [3, 4]), |
| # Lossy DECIMAL scale-down: (10,4) -> (10,2) truncates the extra |
| # scale rather than raising. |
| ("decimal_scale_down", |
| pa.decimal128(10, 4), pa.decimal128(10, 2), 'DECIMAL(10, 2)', |
| [decimal.Decimal('1.2345'), decimal.Decimal('4.5678')], |
| [decimal.Decimal('1.23'), decimal.Decimal('4.56')], |
| [decimal.Decimal('7.89'), decimal.Decimal('0.12')]), |
| ] |
| |
| for (name, old_type, new_type, new_type_str, |
| write_vals, old_read_vals, new_vals) in cases: |
| with self.subTest(case=name): |
| table_name = f'default.promo_{name}' |
| old_schema = pa.schema([('k', pa.int64()), ('v', old_type)]) |
| self.catalog.create_table( |
| table_name, Schema.from_pyarrow_schema(old_schema), False) |
| |
| # Write under the original schema (file stamped schema_id 0). |
| table = self.catalog.get_table(table_name) |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| table_write.write_arrow(pa.Table.from_pydict( |
| {'k': [1, 2], 'v': write_vals}, schema=old_schema)) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # Widen column v through the catalog (new schema_id 1). |
| self.catalog.alter_table( |
| table_name, |
| [SchemaChange.update_column_type( |
| 'v', AtomicType(new_type_str))], |
| False) |
| |
| # Write under the promoted schema (file stamped schema_id 1). |
| table = self.catalog.get_table(table_name) |
| new_schema = pa.schema([('k', pa.int64()), ('v', new_type)]) |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| table_write.write_arrow(pa.Table.from_pydict( |
| {'k': [3, 4], 'v': new_vals}, schema=new_schema)) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| # Plain full-table read spanning both schema versions. |
| read_builder = table.new_read_builder() |
| actual = read_builder.new_read().to_arrow( |
| self._scan_table(read_builder)) |
| expected = pa.Table.from_pydict( |
| {'k': [1, 2, 3, 4], 'v': old_read_vals + new_vals}, |
| schema=new_schema) |
| self.assertEqual(expected, actual) |
| |
| def test_schema_evolution_type_lossy_old_file_only(self): |
| # Reading ONLY old-schema files after a lossy type change (no |
| # newer-schema file in the splits). The output type must equal the |
| # current read schema regardless of which files the read spans, and the |
| # conversion must truncate to match Java CastExecutors rather than |
| # raise. (A previous fix that relied on pyarrow's safe cast crashed |
| # here on lossy evolutions.) |
| import decimal |
| |
| cases = [ |
| ("scale_down", |
| pa.decimal128(10, 4), pa.decimal128(10, 2), 'DECIMAL(10, 2)', |
| [decimal.Decimal('1.2345'), decimal.Decimal('4.5678')], |
| [decimal.Decimal('1.23'), decimal.Decimal('4.56')]), |
| ("double_to_int", pa.float64(), pa.int32(), 'INT', |
| [1.2, 2.8], [1, 2]), |
| ] |
| |
| for name, old_type, new_type, new_type_str, write_vals, read_vals \ |
| in cases: |
| with self.subTest(case=name): |
| table_name = f'default.lossy_old_only_{name}' |
| old_schema = pa.schema([('k', pa.int64()), ('v', old_type)]) |
| self.catalog.create_table( |
| table_name, Schema.from_pyarrow_schema(old_schema), False) |
| |
| # Write under the original schema, then change the type. No |
| # write happens afterwards, so the read sees only this file. |
| table = self.catalog.get_table(table_name) |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| table_write.write_arrow(pa.Table.from_pydict( |
| {'k': [1, 2], 'v': write_vals}, schema=old_schema)) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| self.catalog.alter_table( |
| table_name, |
| [SchemaChange.update_column_type( |
| 'v', AtomicType(new_type_str))], |
| False) |
| |
| table = self.catalog.get_table(table_name) |
| new_schema = pa.schema([('k', pa.int64()), ('v', new_type)]) |
| read_builder = table.new_read_builder() |
| actual = read_builder.new_read().to_arrow( |
| self._scan_table(read_builder)) |
| expected = pa.Table.from_pydict( |
| {'k': [1, 2], 'v': read_vals}, schema=new_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()) |
| ]) |
| options = {'metadata.stats-mode': 'full'} |
| schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt'], options=options) |
| 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().file_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.file_scanner.manifest_list_manager.read_all(latest_snapshot) |
| entries = new_scan.file_scanner.read_manifest_entries(manifest_files) |
| self.assertEqual(1, len(entries)) # verify scan filter success for schema evolution |
| |
| @pytest.mark.python_plan |
| 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()) |
| ]) |
| options = {'metadata.stats-mode': 'full'} |
| schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt'], options=options) |
| 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'], options=options) |
| 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.file_io.delete_quietly(table2.table_path + "/schema/schema-0") |
| 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) |
| |
| # ------------------------------------------------------------------ |
| # Public-API end-to-end evolution cases (create -> write -> alter_table |
| # -> write -> read), the范式 promoted by |
| # test_schema_evolution_type_promotion_unpartitioned. Unlike the older |
| # white-box cases above (manual schema-N files + file.schema_id), these |
| # drive evolution purely through catalog.alter_table. |
| # ------------------------------------------------------------------ |
| |
| def _write(self, table, pa_table): |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| def _read_rows(self, table, sort_key='k'): |
| read_builder = table.new_read_builder() |
| rows = read_builder.new_read().to_arrow( |
| self._scan_table(read_builder)).to_pylist() |
| return sorted(rows, key=lambda r: (r[sort_key] is None, r[sort_key])) |
| |
| def _read_arrow(self, table, projection=None): |
| read_builder = table.new_read_builder() |
| if projection is not None: |
| read_builder = read_builder.with_projection(projection) |
| return read_builder.new_read().to_arrow(self._scan_table(read_builder)) |
| |
| def test_evolution_drop_column_then_read(self): |
| # A5: a column dropped after old data was written must disappear from |
| # the read; the old file's value for it is discarded, not surfaced. |
| name = 'default.evo_drop' |
| s0 = pa.schema([('k', pa.int64()), ('v', pa.string()), ('w', pa.string())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1, 2], 'v': ['a', 'b'], 'w': ['p', 'q']}, schema=s0)) |
| |
| self.catalog.alter_table(name, [SchemaChange.drop_column('v')], False) |
| table = self.catalog.get_table(name) |
| s1 = pa.schema([('k', pa.int64()), ('w', pa.string())]) |
| self._write(table, pa.Table.from_pydict({'k': [3], 'w': ['r']}, schema=s1)) |
| |
| self.assertEqual(self._read_rows(table), [ |
| {'k': 1, 'w': 'p'}, {'k': 2, 'w': 'q'}, {'k': 3, 'w': 'r'}]) |
| |
| def test_evolution_rename_column_then_read(self): |
| # A6: after rename the old file (written under the old name) and the |
| # new file (new name) must read back as the SAME logical column, |
| # matched by field id, not by name. |
| name = 'default.evo_rename' |
| s0 = pa.schema([('k', pa.int64()), ('v', pa.string())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1, 2], 'v': ['a', 'b']}, schema=s0)) |
| |
| self.catalog.alter_table( |
| name, [SchemaChange.rename_column('v', 'renamed')], False) |
| table = self.catalog.get_table(name) |
| s1 = pa.schema([('k', pa.int64()), ('renamed', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [3], 'renamed': ['c']}, schema=s1)) |
| |
| self.assertEqual(self._read_rows(table), [ |
| {'k': 1, 'renamed': 'a'}, |
| {'k': 2, 'renamed': 'b'}, |
| {'k': 3, 'renamed': 'c'}]) |
| |
| def test_evolution_column_position_then_read(self): |
| # A7: moving a column to FIRST must reorder the read schema; the old |
| # file (original order) must be remapped so values stay with their |
| # column. |
| name = 'default.evo_position' |
| s0 = pa.schema([('k', pa.int64()), ('a', pa.string()), ('b', pa.string())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1], 'a': ['a1'], 'b': ['b1']}, schema=s0)) |
| |
| self.catalog.alter_table( |
| name, [SchemaChange.update_column_position(Move.first('b'))], False) |
| table = self.catalog.get_table(name) |
| s1 = pa.schema([('b', pa.string()), ('k', pa.int64()), ('a', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'b': ['b2'], 'k': [2], 'a': ['a2']}, schema=s1)) |
| |
| arrow = self._read_arrow(table) |
| self.assertEqual(arrow.column_names, ['b', 'k', 'a']) |
| self.assertEqual(self._read_rows(table), [ |
| {'b': 'b1', 'k': 1, 'a': 'a1'}, |
| {'b': 'b2', 'k': 2, 'a': 'a2'}]) |
| |
| def test_evolution_multi_version_chain(self): |
| # A8: chain schema versions (add + type promotion), writing one batch |
| # per version then reading across all of them. Rename across versions |
| # is covered on its own by A6; this case isolates the add+promotion |
| # chain. |
| name = 'default.evo_chain' |
| s0 = pa.schema([('k', pa.int64()), ('a', pa.string()), ('n', pa.int32())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1], 'a': ['a1'], 'n': [10]}, schema=s0)) |
| |
| # v1: add column b |
| self.catalog.alter_table( |
| name, [SchemaChange.add_column('b', AtomicType('STRING'))], False) |
| table = self.catalog.get_table(name) |
| s1 = pa.schema([('k', pa.int64()), ('a', pa.string()), |
| ('n', pa.int32()), ('b', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [2], 'a': ['a2'], 'n': [20], 'b': ['b2']}, schema=s1)) |
| |
| # v2: promote n int32 -> int64 |
| self.catalog.alter_table( |
| name, [SchemaChange.update_column_type('n', AtomicType('BIGINT'))], |
| False) |
| table = self.catalog.get_table(name) |
| s2 = pa.schema([('k', pa.int64()), ('a', pa.string()), |
| ('n', pa.int64()), ('b', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [3], 'a': ['a3'], 'n': [30], 'b': ['b3']}, schema=s2)) |
| |
| # v3: add column d |
| self.catalog.alter_table( |
| name, [SchemaChange.add_column('d', AtomicType('STRING'))], False) |
| table = self.catalog.get_table(name) |
| s3 = pa.schema([('k', pa.int64()), ('a', pa.string()), |
| ('n', pa.int64()), ('b', pa.string()), ('d', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [4], 'a': ['a4'], 'n': [40], 'b': ['b4'], 'd': ['d4']}, |
| schema=s3)) |
| |
| self.assertEqual(self._read_rows(table), [ |
| {'k': 1, 'a': 'a1', 'n': 10, 'b': None, 'd': None}, |
| {'k': 2, 'a': 'a2', 'n': 20, 'b': 'b2', 'd': None}, |
| {'k': 3, 'a': 'a3', 'n': 30, 'b': 'b3', 'd': None}, |
| {'k': 4, 'a': 'a4', 'n': 40, 'b': 'b4', 'd': 'd4'}]) |
| |
| def test_evolution_projection_after_add_column(self): |
| # A10: projecting a subset that includes the newly added column must |
| # return NULL for that column on rows from old files. |
| name = 'default.evo_projection' |
| s0 = pa.schema([('k', pa.int64()), ('a', pa.string())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1, 2], 'a': ['a1', 'a2']}, schema=s0)) |
| |
| self.catalog.alter_table( |
| name, [SchemaChange.add_column('b', AtomicType('STRING'))], False) |
| table = self.catalog.get_table(name) |
| s1 = pa.schema([('k', pa.int64()), ('a', pa.string()), ('b', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [3], 'a': ['a3'], 'b': ['b3']}, schema=s1)) |
| |
| rows = sorted( |
| self._read_arrow(table, projection=['k', 'b']).to_pylist(), |
| key=lambda r: r['k']) |
| self.assertEqual(rows, [ |
| {'k': 1, 'b': None}, {'k': 2, 'b': None}, {'k': 3, 'b': 'b3'}]) |
| |
| def test_evolution_nullability_then_read(self): |
| # A11: relaxing NOT NULL -> nullable must let later NULLs read back |
| # while the old (non-null) file still reads correctly. |
| name = 'default.evo_nullability' |
| s0 = pa.schema([('k', pa.int64()), |
| pa.field('v', pa.string(), nullable=False)]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1, 2], 'v': ['a', 'b']}, schema=s0)) |
| |
| self.catalog.alter_table( |
| name, [SchemaChange.update_column_nullability('v', True)], False) |
| table = self.catalog.get_table(name) |
| s1 = pa.schema([('k', pa.int64()), ('v', pa.string())]) |
| self._write(table, pa.Table.from_pydict({'k': [3], 'v': [None]}, schema=s1)) |
| |
| self.assertEqual(self._read_rows(table), [ |
| {'k': 1, 'v': 'a'}, {'k': 2, 'v': 'b'}, {'k': 3, 'v': None}]) |
| |
| def test_evolution_drop_then_readd_same_name(self): |
| # A13 (name-vs-field-id isolation): dropping v then adding a new v |
| # (fresh field id) must NOT revive the old file's v. Because columns |
| # are aligned by field id, the dropped v's id is gone and the re-added |
| # v's new id is absent from old files, so old rows read NULL. Covers |
| # same-type and changed-type re-add. |
| for label, readd_paimon, readd_arrow, new_val in [ |
| ('same_type', 'STRING', pa.string(), ['new3']), |
| ('changed_type', 'INT', pa.int32(), [99]), |
| ]: |
| with self.subTest(case=label): |
| name = 'default.evo_drop_readd_{}'.format(label) |
| s0 = pa.schema([('k', pa.int64()), ('v', pa.string())]) |
| self.catalog.create_table( |
| name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1, 2], 'v': ['old1', 'old2']}, schema=s0)) |
| |
| self.catalog.alter_table( |
| name, [SchemaChange.drop_column('v')], False) |
| self.catalog.alter_table( |
| name, [SchemaChange.add_column('v', AtomicType(readd_paimon))], |
| False) |
| table = self.catalog.get_table(name) |
| s2 = pa.schema([('k', pa.int64()), ('v', readd_arrow)]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [3], 'v': new_val}, schema=s2)) |
| |
| self.assertEqual(self._read_rows(table), [ |
| {'k': 1, 'v': None}, |
| {'k': 2, 'v': None}, |
| {'k': 3, 'v': new_val[0]}]) |
| |
| def test_evolution_rename_then_add_same_name(self): |
| # A14 (name-vs-field-id isolation): rename a->b, then add a new a |
| # (fresh field id). Aligned by field id, reading old files b (old id) |
| # carries the old a-data and the new a (new id) reads NULL -- a |
| # name-based alignment would instead feed the old physical column 'a' |
| # to the new a and leave b empty. |
| name = 'default.evo_rename_readd' |
| s0 = pa.schema([('k', pa.int64()), ('a', pa.string())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1, 2], 'a': ['old1', 'old2']}, schema=s0)) |
| |
| self.catalog.alter_table( |
| name, [SchemaChange.rename_column('a', 'b')], False) |
| self.catalog.alter_table( |
| name, [SchemaChange.add_column('a', AtomicType('STRING'))], False) |
| table = self.catalog.get_table(name) |
| s2 = pa.schema([('k', pa.int64()), ('b', pa.string()), ('a', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [3], 'b': ['B3'], 'a': ['A3']}, schema=s2)) |
| |
| self.assertEqual(self._read_rows(table), [ |
| {'k': 1, 'b': 'old1', 'a': None}, |
| {'k': 2, 'b': 'old2', 'a': None}, |
| {'k': 3, 'b': 'B3', 'a': 'A3'}]) |
| |
| def test_evolution_column_swap_direct_rejected(self): |
| # A15a: a one-shot name swap is rejected -- renaming a->b while b |
| # still exists collides on the existing name. (Constraint check, this |
| # behaviour is correct.) |
| name = 'default.evo_swap_direct' |
| s0 = pa.schema([('k', pa.int64()), ('a', pa.string()), ('b', pa.string())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| # alter_table wraps the underlying ColumnAlreadyExistException in a |
| # RuntimeError (filesystem_catalog catch-all), so match on that. |
| with self.assertRaises(RuntimeError) as cm: |
| self.catalog.alter_table(name, [ |
| SchemaChange.rename_column('a', 'b'), |
| SchemaChange.rename_column('b', 'a'), |
| ], False) |
| self.assertIn('already exists', str(cm.exception)) |
| |
| def test_evolution_column_swap_via_temp_name(self): |
| # A15b (name-vs-field-id isolation): a 3-step swap via a temp name |
| # keeps field ids stable, so on read the old data follows the id, not |
| # the name -- old 'a1' (id of a) ends up under column b, old 'b1' |
| # under column a. A name-based alignment would drop/misalign both. |
| name = 'default.evo_swap_temp' |
| s0 = pa.schema([('k', pa.int64()), ('a', pa.string()), ('b', pa.string())]) |
| self.catalog.create_table(name, Schema.from_pyarrow_schema(s0), False) |
| table = self.catalog.get_table(name) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [1], 'a': ['a1'], 'b': ['b1']}, schema=s0)) |
| |
| self.catalog.alter_table(name, [ |
| SchemaChange.rename_column('a', '__tmp'), |
| SchemaChange.rename_column('b', 'a'), |
| SchemaChange.rename_column('__tmp', 'b'), |
| ], False) |
| table = self.catalog.get_table(name) |
| s1 = pa.schema([('k', pa.int64()), ('b', pa.string()), ('a', pa.string())]) |
| self._write(table, pa.Table.from_pydict( |
| {'k': [2], 'b': ['B2'], 'a': ['A2']}, schema=s1)) |
| |
| self.assertEqual(self._read_rows(table), [ |
| {'k': 1, 'b': 'a1', 'a': 'b1'}, |
| {'k': 2, 'b': 'B2', 'a': 'A2'}]) |
| |
| 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 |
