From 3030ee938a598899004efad75679e94afa3e8a5e Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Thu, 20 Aug 2026 16:32:19 -0700 Subject: [PATCH 1/4] Fix upsert after schema evolution --- pyiceberg/table/__init__.py | 19 ++++++------- tests/table/test_upsert.py | 54 +++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 9 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index bb879dfbce..baf7e28109 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -892,7 +892,7 @@ def upsert( except ModuleNotFoundError as e: raise ModuleNotFoundError("For writes PyArrow needs to be installed") from e - from pyiceberg.io.pyarrow import expression_to_pyarrow + from pyiceberg.io.pyarrow import ArrowScan, expression_to_pyarrow from pyiceberg.table import upsert_util if join_cols is None: @@ -926,19 +926,20 @@ def upsert( # get list of rows that exist so we don't have to load the entire target table matched_predicate = upsert_util.create_match_filter(df, join_cols) - # We must use Transaction.table_metadata for the scan. This includes all uncommitted - but relevant - changes. + matched_iceberg_file_scan = self._scan(row_filter=matched_predicate, case_sensitive=case_sensitive) - matched_iceberg_record_batches_scan = DataScan( + if branch in self.table_metadata.refs: + matched_iceberg_file_scan = matched_iceberg_file_scan.use_ref(branch) + + # The target branch determines which files to read; the transaction schema determines how to project their rows. + # These can differ because schema updates do not create snapshots. + matched_iceberg_record_batches = ArrowScan( table_metadata=self.table_metadata, io=self._table.io, + projected_schema=self.table_metadata.schema(), row_filter=matched_predicate, case_sensitive=case_sensitive, - ) - - if branch in self.table_metadata.refs: - matched_iceberg_record_batches_scan = matched_iceberg_record_batches_scan.use_ref(branch) - - matched_iceberg_record_batches = matched_iceberg_record_batches_scan.to_arrow_batch_reader() + ).to_record_batches(matched_iceberg_file_scan.plan_files()) batches_to_overwrite = [] overwrite_predicates = [] diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..196a19022b 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -30,6 +30,7 @@ from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema from pyiceberg.table import Table, UpsertResult +from pyiceberg.table.refs import MAIN_BRANCH from pyiceberg.table.snapshots import Operation from pyiceberg.table.upsert_util import create_match_filter from pyiceberg.transforms import DayTransform @@ -391,6 +392,59 @@ def test_upsert_with_identifier_fields(catalog: Catalog) -> None: assert [snap.summary.operation for snap in tbl.snapshots() if snap.summary is not None] == expected_operations +@pytest.mark.parametrize("branch", [MAIN_BRANCH, "test_branch"]) +def test_upsert_after_schema_evolution(catalog: Catalog, branch: str) -> None: + identifier = "default.test_upsert_after_schema_evolution" + schema = Schema( + NestedField(1, "city", StringType(), required=True), + NestedField(2, "population", IntegerType(), required=True), + identifier_field_ids=[1], + ) + tbl = catalog.create_table(identifier, schema=schema) + initial_arrow_schema = pa.schema( + [ + pa.field("city", pa.string(), nullable=False), + pa.field("population", pa.int32(), nullable=False), + ] + ) + tbl.append(pa.Table.from_pylist([{"city": "Amsterdam", "population": 921402}], schema=initial_arrow_schema)) + + initial_snapshot = tbl.current_snapshot() + assert initial_snapshot is not None + + if branch != MAIN_BRANCH: + tbl.manage_snapshots().create_branch(snapshot_id=initial_snapshot.snapshot_id, branch_name=branch).commit() + tbl.delete("city == 'Amsterdam'") + + snapshot_before_schema_update = tbl.current_snapshot() + assert snapshot_before_schema_update is not None + + with tbl.update_schema() as update: + update.add_column("country", StringType()) + + assert tbl.schema().schema_id != snapshot_before_schema_update.schema_id + assert tbl.metadata.current_snapshot_id == snapshot_before_schema_update.snapshot_id + + evolved_arrow_schema = pa.schema( + [ + pa.field("city", pa.string(), nullable=False), + pa.field("population", pa.int32(), nullable=False), + pa.field("country", pa.string()), + ] + ) + source = pa.Table.from_pylist( + [{"city": "Amsterdam", "population": 934927, "country": "Netherlands"}], schema=evolved_arrow_schema + ) + + result = tbl.upsert(source, branch=branch) + + assert_upsert_result(result, expected_updated=1, expected_inserted=0) + result_scan = tbl.scan() if branch == MAIN_BRANCH else tbl.scan().use_ref(branch) + assert result_scan.to_arrow().to_pylist() == [{"city": "Amsterdam", "population": 934927, "country": "Netherlands"}] + if branch != MAIN_BRANCH: + assert tbl.scan().to_arrow().to_pylist() == [] + + def test_upsert_into_empty_table(catalog: Catalog) -> None: identifier = "default.test_upsert_into_empty_table" _drop_table(catalog, identifier) From 7cba60b80acdec76c1e3d7792bd1bd51e6bfd432 Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Thu, 20 Aug 2026 16:39:14 -0700 Subject: [PATCH 2/4] Split schema evolution upsert tests --- tests/table/test_upsert.py | 60 +++++++++++++++++++++++++++++++------- 1 file changed, 50 insertions(+), 10 deletions(-) diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 196a19022b..39333f65ab 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -30,7 +30,6 @@ from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema from pyiceberg.table import Table, UpsertResult -from pyiceberg.table.refs import MAIN_BRANCH from pyiceberg.table.snapshots import Operation from pyiceberg.table.upsert_util import create_match_filter from pyiceberg.transforms import DayTransform @@ -392,8 +391,7 @@ def test_upsert_with_identifier_fields(catalog: Catalog) -> None: assert [snap.summary.operation for snap in tbl.snapshots() if snap.summary is not None] == expected_operations -@pytest.mark.parametrize("branch", [MAIN_BRANCH, "test_branch"]) -def test_upsert_after_schema_evolution(catalog: Catalog, branch: str) -> None: +def test_upsert_after_schema_evolution(catalog: Catalog) -> None: identifier = "default.test_upsert_after_schema_evolution" schema = Schema( NestedField(1, "city", StringType(), required=True), @@ -409,12 +407,54 @@ def test_upsert_after_schema_evolution(catalog: Catalog, branch: str) -> None: ) tbl.append(pa.Table.from_pylist([{"city": "Amsterdam", "population": 921402}], schema=initial_arrow_schema)) + snapshot_before_schema_update = tbl.current_snapshot() + assert snapshot_before_schema_update is not None + + with tbl.update_schema() as update: + update.add_column("country", StringType()) + + assert tbl.schema().schema_id != snapshot_before_schema_update.schema_id + assert tbl.metadata.current_snapshot_id == snapshot_before_schema_update.snapshot_id + + evolved_arrow_schema = pa.schema( + [ + pa.field("city", pa.string(), nullable=False), + pa.field("population", pa.int32(), nullable=False), + pa.field("country", pa.string()), + ] + ) + source = pa.Table.from_pylist( + [{"city": "Amsterdam", "population": 934927, "country": "Netherlands"}], schema=evolved_arrow_schema + ) + + result = tbl.upsert(source) + + assert_upsert_result(result, expected_updated=1, expected_inserted=0) + assert tbl.scan().to_arrow().to_pylist() == [{"city": "Amsterdam", "population": 934927, "country": "Netherlands"}] + + +def test_upsert_to_branch_after_schema_evolution(catalog: Catalog) -> None: + identifier = "default.test_upsert_to_branch_after_schema_evolution" + schema = Schema( + NestedField(1, "city", StringType(), required=True), + NestedField(2, "population", IntegerType(), required=True), + identifier_field_ids=[1], + ) + tbl = catalog.create_table(identifier, schema=schema) + initial_arrow_schema = pa.schema( + [ + pa.field("city", pa.string(), nullable=False), + pa.field("population", pa.int32(), nullable=False), + ] + ) + tbl.append(pa.Table.from_pylist([{"city": "Amsterdam", "population": 921402}], schema=initial_arrow_schema)) + initial_snapshot = tbl.current_snapshot() assert initial_snapshot is not None - if branch != MAIN_BRANCH: - tbl.manage_snapshots().create_branch(snapshot_id=initial_snapshot.snapshot_id, branch_name=branch).commit() - tbl.delete("city == 'Amsterdam'") + branch = "test_branch" + tbl.manage_snapshots().create_branch(snapshot_id=initial_snapshot.snapshot_id, branch_name=branch).commit() + tbl.delete("city == 'Amsterdam'") snapshot_before_schema_update = tbl.current_snapshot() assert snapshot_before_schema_update is not None @@ -439,10 +479,10 @@ def test_upsert_after_schema_evolution(catalog: Catalog, branch: str) -> None: result = tbl.upsert(source, branch=branch) assert_upsert_result(result, expected_updated=1, expected_inserted=0) - result_scan = tbl.scan() if branch == MAIN_BRANCH else tbl.scan().use_ref(branch) - assert result_scan.to_arrow().to_pylist() == [{"city": "Amsterdam", "population": 934927, "country": "Netherlands"}] - if branch != MAIN_BRANCH: - assert tbl.scan().to_arrow().to_pylist() == [] + assert tbl.scan().use_ref(branch).to_arrow().to_pylist() == [ + {"city": "Amsterdam", "population": 934927, "country": "Netherlands"} + ] + assert tbl.scan().to_arrow().to_pylist() == [] def test_upsert_into_empty_table(catalog: Catalog) -> None: From fa4a0af5a35e78024f7495448fb9edbf8003a171 Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Thu, 20 Aug 2026 16:45:22 -0700 Subject: [PATCH 3/4] Move branch selection into transaction scan --- pyiceberg/table/__init__.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index baf7e28109..f52dbb3940 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -314,11 +314,19 @@ def _apply( return self - def _scan(self, row_filter: str | BooleanExpression = ALWAYS_TRUE, case_sensitive: bool = True) -> DataScan: - """Minimal data scan of the table with the current state of the transaction.""" - return DataScan( + def _scan( + self, + row_filter: str | BooleanExpression = ALWAYS_TRUE, + case_sensitive: bool = True, + branch: str | None = None, + ) -> DataScan: + """Minimal data scan of the current transaction state, optionally scoped to a branch.""" + scan = DataScan( table_metadata=self.table_metadata, io=self._table.io, row_filter=row_filter, case_sensitive=case_sensitive ) + if branch in self.table_metadata.refs: + return scan.use_ref(branch) + return scan def upgrade_table_version(self, format_version: TableVersion) -> Transaction: """Set the table to a certain version. @@ -778,9 +786,7 @@ def delete( bound_delete_filter = bind(self.table_metadata.schema(), delete_filter, case_sensitive) preserve_row_filter = _expression_to_complementary_pyarrow(bound_delete_filter, self.table_metadata.schema()) - file_scan = self._scan(row_filter=delete_filter, case_sensitive=case_sensitive) - if branch is not None: - file_scan = file_scan.use_ref(branch) + file_scan = self._scan(row_filter=delete_filter, case_sensitive=case_sensitive, branch=branch) files = file_scan.plan_files() commit_uuid = uuid.uuid4() @@ -926,10 +932,7 @@ def upsert( # get list of rows that exist so we don't have to load the entire target table matched_predicate = upsert_util.create_match_filter(df, join_cols) - matched_iceberg_file_scan = self._scan(row_filter=matched_predicate, case_sensitive=case_sensitive) - - if branch in self.table_metadata.refs: - matched_iceberg_file_scan = matched_iceberg_file_scan.use_ref(branch) + matched_iceberg_file_scan = self._scan(row_filter=matched_predicate, case_sensitive=case_sensitive, branch=branch) # The target branch determines which files to read; the transaction schema determines how to project their rows. # These can differ because schema updates do not create snapshots. From 6e5522b5d7b878c0634e21a214960f7d49b9495e Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Thu, 20 Aug 2026 16:54:34 -0700 Subject: [PATCH 4/4] Preserve batch reader behavior in upsert --- pyiceberg/table/__init__.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index f52dbb3940..5dcbac4915 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -898,7 +898,7 @@ def upsert( except ModuleNotFoundError as e: raise ModuleNotFoundError("For writes PyArrow needs to be installed") from e - from pyiceberg.io.pyarrow import ArrowScan, expression_to_pyarrow + from pyiceberg.io.pyarrow import expression_to_pyarrow from pyiceberg.table import upsert_util if join_cols is None: @@ -934,15 +934,13 @@ def upsert( matched_iceberg_file_scan = self._scan(row_filter=matched_predicate, case_sensitive=case_sensitive, branch=branch) - # The target branch determines which files to read; the transaction schema determines how to project their rows. - # These can differ because schema updates do not create snapshots. - matched_iceberg_record_batches = ArrowScan( - table_metadata=self.table_metadata, - io=self._table.io, - projected_schema=self.table_metadata.schema(), - row_filter=matched_predicate, - case_sensitive=case_sensitive, - ).to_record_batches(matched_iceberg_file_scan.plan_files()) + # Plan files from the target branch, but read them using the transaction's current schema. + # A schema update does not create a snapshot, so the branch snapshot may use an older schema. + matched_iceberg_record_batches = _to_arrow_batch_reader_via_file_scan_tasks( + matched_iceberg_file_scan, + self.table_metadata.schema(), + matched_iceberg_file_scan.plan_files(), + ) batches_to_overwrite = [] overwrite_predicates = []