Fix upsert after schema evolution

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,20 +932,16 @@ 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, branch=branch)
 
-        matched_iceberg_record_batches_scan = DataScan(
-            table_metadata=self.table_metadata,
-            io=self._table.io,
-            row_filter=matched_predicate,
-            case_sensitive=case_sensitive,
+        # 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(),
         )
 
-        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()
-
         batches_to_overwrite = []
         overwrite_predicates = []
         rows_to_insert = df

tests/table/test_upsert.py

@@ -391,6 +391,100 @@ 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
 
 
+def test_upsert_after_schema_evolution(catalog: Catalog) -> 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))
+
+    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
+
+    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
+
+    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)
+    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:
     identifier = "default.test_upsert_into_empty_table"
     _drop_table(catalog, identifier)