diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index bb879dfbce..5dcbac4915 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,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 diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..39333f65ab 100644 --- a/tests/table/test_upsert.py +++ b/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)