[Python] Resolve time-travel schemas from the selected branch

paimon-python/pypaimon/table/file_store_table.py

@@ -532,10 +532,6 @@ def _copy(self, options: dict, resolve_time_travel: bool) -> 'FileStoreTable':
             key in options for key in SCAN_KEYS + [
                 CoreOptions.SCAN_MODE.key(), CoreOptions.BRANCH.key(),
                 CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key()])
-        if resolve_time_travel and not preserve_snapshot:
-            time_travel_schema = self._try_time_travel(Options(new_options), set(applied_options))
-            if time_travel_schema is not None:
-                new_table_schema = time_travel_schema
 
         # Re-encode the branch into the identifier when the option changes, so
         # current_branch() and any catalog-routed snapshot commit see the
@@ -554,6 +550,13 @@ def _copy(self, options: dict, resolve_time_travel: bool) -> 'FileStoreTable':
 
         new_table = FileStoreTable(self.file_io, new_identifier, self.table_path,
                                    new_table_schema, catalog_env)
+        if resolve_time_travel and not preserve_snapshot:
+            # Resolve snapshots, tags and schemas in the target branch.
+            time_travel_schema = new_table._try_time_travel(Options(new_options), set(applied_options))
+            if time_travel_schema is not None:
+                new_table = FileStoreTable(self.file_io, new_identifier, self.table_path,
+                                           time_travel_schema, catalog_env)
+
         new_table._applied_dynamic_options = applied_options
         if preserve_snapshot:
             new_table._read_snapshot = self._read_snapshot

paimon-python/pypaimon/tests/table/simple_table_test.py

@@ -24,7 +24,8 @@
 
 from pypaimon import CatalogFactory, Schema
 from pypaimon.common.options.core_options import CoreOptions
-from pypaimon.schema.schema_manager import SchemaManager
+from pypaimon.schema.data_types import AtomicType
+from pypaimon.schema.schema_change import SchemaChange
 
 
 class SimpleTableTest(unittest.TestCase):
@@ -492,107 +493,59 @@ def test_tag_rename_and_scan(self):
         self.assertEqual(result_sorted.column('v').to_pylist(), expected.column('v').to_pylist())
 
     def test_schema_evolution_tag_read(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_0', schema, False)
-        table1 = self.catalog.get_table('default.test_0')
-        write_builder = table1.new_batch_write_builder()
-        # write 1
-        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()
-        # write 2
-        table_write = write_builder.new_write()
-        table_commit = write_builder.new_commit()
-        data1 = {
-            'user_id': [11, 22, 33, 44],
-            '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_1', schema2, False)
-        table2 = self.catalog.get_table('default.test_1')
-        write_builder = table2.new_batch_write_builder()
-        # write 1
-        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 2
+        name = 'default.test_schema_evolution'
+        self.catalog.create_table(name, schema, False)
+        table = self.catalog.get_table(name)
+        expected = pa.Table.from_pydict({
+            'user_id': [5, 6, 7, 8, 55, 66, 77, 88],
+            'item_id': [1005, 1006, 1007, 1008, 1005, 1006, 1007, 1008],
+            'dt': ["p2", "p1", "p2", "p2", "p2", "p1", "p2", "p2"]
+        }, schema=pa_schema)
+        write_builder = table.new_batch_write_builder()
+        for data in (expected.slice(0, 4), expected.slice(4)):
+            table_write = write_builder.new_write()
+            table_commit = write_builder.new_commit()
+            try:
+                table_write.write_arrow(data)
+                table_commit.commit(table_write.prepare_commit())
+            finally:
+                table_write.close()
+                table_commit.close()
+        table.create_tag("tag2", snapshot_id=2)
+
+        # Evolve the same table and write newer data beyond the tagged snapshot.
+        self.catalog.alter_table(name, [SchemaChange.add_column('behavior', AtomicType('STRING'))])
+        table = self.catalog.get_table(name)
+        self.assertEqual(table.table_schema.id, 1)
+        write_builder = table.new_batch_write_builder()
         table_write = write_builder.new_write()
         table_commit = write_builder.new_commit()
-        data2 = {
-            'user_id': [55, 66, 77, 88],
-            '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()
+        try:
+            table_write.write_arrow(pa.Table.from_pydict({
+                'user_id': [99],
+                'item_id': [1009],
+                'dt': ['p3'],
+                'behavior': ['new'],
+            }, schema=pa_schema.append(pa.field('behavior', pa.string()))))
+            table_commit.commit(table_write.prepare_commit())
+        finally:
+            table_write.close()
+            table_commit.close()
 
-        # Create tag at snapshot 2
-        table1.create_tag("tag2", snapshot_id=2)
-        table2.create_tag("tag2", snapshot_id=2)
-        # When table2 read tag2, it will access table1's schema
-        table2.schema_manager = SchemaManager(table2.file_io, table1.table_path)
-        table_with_tag = table2.copy({CoreOptions.SCAN_TAG_NAME.key(): "tag2"})
+        table_with_tag = table.copy({CoreOptions.SCAN_TAG_NAME.key(): "tag2"})
+        self.assertEqual(table_with_tag.table_schema.id, 0)
         read_builder = table_with_tag.new_read_builder()
         table_scan = read_builder.new_scan()
         table_read = read_builder.new_read()
 
         result = table_read.to_arrow(table_scan.plan().splits())
 
-        self.assertEqual(result.num_rows, 8)
-
-        expected_schema = pa.schema([
-            ('user_id', pa.int64()),
-            ('item_id', pa.int64()),
-            ('dt', pa.string())
-        ])
-        expected = pa.Table.from_pydict({
-            'user_id': [5, 6, 7, 8, 55, 66, 77, 88],
-            'item_id': [1005, 1006, 1007, 1008, 1005, 1006, 1007, 1008],
-            'dt': ["p2", "p1", "p2", "p2", "p2", "p1", "p2", "p2"]
-        }, schema=expected_schema)
-
         self.assertEqual(expected, result.sort_by('user_id'))
 
     def test_table_rollback_to_snapshot(self):

paimon-python/pypaimon/tests/table/time_travel_schema_evolution_test.py

@@ -20,6 +20,7 @@
 import pytest
 
 from pypaimon import CatalogFactory, Schema
+from pypaimon.read.datasource.split_provider import CatalogSplitProvider
 from pypaimon.schema.data_types import ArrayType, AtomicType, DataField, PyarrowFieldParser
 from pypaimon.schema.schema_change import SchemaChange
 from pypaimon.table.row.blob import BlobDescriptor, BlobViewStruct
@@ -64,6 +65,103 @@ def _read(table):
         builder = table.new_read_builder()
         return builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist()
 
+    def _create_branched_table(self, fields, row):
+        table = self._create(fields, {
+            'row-tracking.enabled': 'false',
+            'data-evolution.enabled': 'false',
+            'scan.native-plan.enabled': 'false',
+            'read.native.enabled': 'false',
+            'write.native.enabled': 'false',
+        })
+        self._write(table, row)
+        table.create_tag('base', 1)
+        self.catalog.create_branch('test.t', 'dev', tag_name='base')
+        return table
+
+    @pytest.mark.parametrize('evolve_main', [False, True])
+    @pytest.mark.parametrize('selector', ['scan.snapshot-id', 'scan.tag-name'])
+    def test_branch_time_travel_uses_target_schema(self, selector, evolve_main):
+        self._create_branched_table([
+            DataField(0, 'id', AtomicType('INT')),
+            DataField(1, 'payload', AtomicType('STRING')),
+        ], {'id': 1, 'payload': 'base'})
+        if evolve_main:
+            self.catalog.alter_table('test.t', [
+                SchemaChange.add_column('source', AtomicType('STRING'))])
+        self.catalog.alter_table('test.t$branch_dev', [
+            SchemaChange.add_column('quality_score', AtomicType('DOUBLE'))])
+        main = self.catalog.get_table('test.t')
+        dev = self.catalog.get_table('test.t$branch_dev')
+        main_row = {'id': 2, 'payload': 'main'}
+        if evolve_main:
+            main_row['source'] = 'web'
+            assert main.table_schema.id == dev.table_schema.id
+        self._write(main, main_row)
+        self._write(dev, {'id': 3, 'payload': 'dev', 'quality_score': 0.9})
+        dev.create_tag('training_v1', 2)
+        self.catalog.alter_table('test.t$branch_dev', [
+            SchemaChange.add_column('later_column', AtomicType('STRING'))])
+        self._write(self.catalog.get_table('test.t$branch_dev'), {
+            'id': 4, 'payload': 'later', 'quality_score': 0.8, 'later_column': 'new'})
+
+        target = {selector: '2' if selector == 'scan.snapshot-id' else 'training_v1'}
+        historical = main.copy({'branch': 'dev', **target})
+        two_step = main.copy({'branch': 'dev'}).copy(target)
+        expected = [
+            {'id': 1, 'payload': 'base', 'quality_score': None},
+            {'id': 3, 'payload': 'dev', 'quality_score': 0.9},
+        ]
+        assert historical.field_names == two_step.field_names == ['id', 'payload', 'quality_score']
+        assert sorted(self._read(historical), key=lambda row: row['id']) == expected
+
+        provider = CatalogSplitProvider(
+            table_identifier='test.t',
+            catalog_options={'warehouse': str(self.root / 'warehouse')},
+            snapshot_id=2 if selector == 'scan.snapshot-id' else None,
+            tag_name='training_v1' if selector == 'scan.tag-name' else None,
+            dynamic_options={'branch': 'dev'},
+        )
+        assert [field.name for field in provider.read_type()] == historical.field_names
+        read = provider.table().new_read_builder().new_read()
+        assert sorted(read.to_arrow(provider.splits()).to_pylist(), key=lambda row: row['id']) == expected
+
+    def test_branch_time_travel_preserves_bigint_values(self):
+        main = self._create_branched_table([
+            DataField(0, 'id', AtomicType('INT')),
+            DataField(1, 'count', AtomicType('INT')),
+        ], {'id': 1, 'count': 1})
+        self.catalog.alter_table('test.t$branch_dev', [
+            SchemaChange.update_column_type('count', AtomicType('BIGINT'))])
+        self._write(main, {'id': 2, 'count': 2})
+        self._write(self.catalog.get_table('test.t$branch_dev'), {'id': 3, 'count': 3000000000})
+
+        historical = main.copy({'branch': 'dev', 'scan.snapshot-id': '2'})
+
+        assert sorted(self._read(historical), key=lambda row: row['id']) == [
+            {'id': 1, 'count': 1}, {'id': 3, 'count': 3000000000}]
+
+    @pytest.mark.parametrize('branch', ['main', None])
+    def test_branch_time_travel_back_to_main(self, branch):
+        self._create_branched_table([
+            DataField(0, 'id', AtomicType('INT')),
+            DataField(1, 'payload', AtomicType('STRING')),
+        ], {'id': 1, 'payload': 'base'})
+        self.catalog.alter_table('test.t', [
+            SchemaChange.add_column('source', AtomicType('STRING'))])
+        main = self.catalog.get_table('test.t')
+        dev = self.catalog.get_table('test.t$branch_dev')
+        self._write(main, {'id': 2, 'payload': 'main', 'source': 'web'})
+        self._write(dev, {'id': 3, 'payload': 'dev'})
+
+        historical = dev.copy({'branch': branch, 'scan.snapshot-id': '2'})
+
+        assert historical.current_branch() == 'main'
+        assert historical.field_names == ['id', 'payload', 'source']
+        assert sorted(self._read(historical), key=lambda row: row['id']) == [
+            {'id': 1, 'payload': 'base', 'source': None},
+            {'id': 2, 'payload': 'main', 'source': 'web'},
+        ]
+
     def test_commit_retry_refreshes_snapshot_schema(self, monkeypatch):
         original = self._create([DataField(0, 'id', AtomicType('INT'))])
         self._write(original, {'id': 1})