From af9c2e2f761439be84d67d45e961eff809fd9e00 Mon Sep 17 00:00:00 2001 From: hedger9487 Date: Wed, 26 Aug 2026 23:15:38 +0800 Subject: [PATCH] Fix upsert after schema evolution (#3105) * Cast only join_cols schema instead of full table schema in get_rows_to_update * Safely handle missing non-key columns in target_table when comparing rows * Add regression tests for upsert after add_column and union_by_name schema evolution * Verify underlying Parquet file replacement and snapshot operations --- pyiceberg/table/upsert_util.py | 15 ++- tests/table/test_upsert.py | 238 ++++++++++++++++++++++++++++++++- 2 files changed, 245 insertions(+), 8 deletions(-) diff --git a/pyiceberg/table/upsert_util.py b/pyiceberg/table/upsert_util.py index 6f32826eb0..59b25f52ae 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -62,8 +62,8 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols """ all_columns = set(source_table.column_names) join_cols_set = set(join_cols) - non_key_cols = list(all_columns - join_cols_set) + target_columns = set(target_table.column_names) if has_duplicate_rows(target_table, join_cols): raise ValueError("Target table has duplicate rows, aborting upsert") @@ -86,19 +86,20 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols ) from None # Step 1: Prepare source index with join keys and a marker index - # Cast to target table schema, so we can do the join + # Cast join columns to target table schema, so we can do the join # See: https://github.com/apache/arrow/issues/37542 + join_schema = pa.schema([target_table.schema.field(col) for col in join_cols]) source_index = ( - source_table.cast(target_table.schema) - .select(join_cols_set) + source_table.select(join_cols) + .cast(join_schema) .append_column(SOURCE_INDEX_COLUMN_NAME, pa.array(range(len(source_table)))) ) # Step 2: Prepare target index with join keys and a marker - target_index = target_table.select(join_cols_set).append_column(TARGET_INDEX_COLUMN_NAME, pa.array(range(len(target_table)))) + target_index = target_table.select(join_cols).append_column(TARGET_INDEX_COLUMN_NAME, pa.array(range(len(target_table)))) # Step 3: Perform an inner join to find which rows from source exist in target - matching_indices = source_index.join(target_index, keys=list(join_cols_set), join_type="inner") + matching_indices = source_index.join(target_index, keys=join_cols, join_type="inner") # Step 4: Compare all rows using Python to_update_indices = [] @@ -112,7 +113,7 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols for key in non_key_cols: source_val = source_row.column(key)[0].as_py() - target_val = target_row.column(key)[0].as_py() + target_val = target_row.column(key)[0].as_py() if key in target_columns else None if source_val != target_val: to_update_indices.append(source_idx) break diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..d45b167a71 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -33,7 +33,7 @@ from pyiceberg.table.snapshots import Operation from pyiceberg.table.upsert_util import create_match_filter from pyiceberg.transforms import DayTransform -from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType +from pyiceberg.types import IntegerType, LongType, NestedField, StringType, StructType, TimestampType from tests.catalog.test_base import InMemoryCatalog @@ -927,3 +927,239 @@ def test_upsert_snapshot_properties(catalog: Catalog) -> None: for snapshot in snapshots[initial_snapshot_count:]: assert snapshot.summary is not None assert snapshot.summary.additional_properties.get("test_prop") == "test_value" + + +def test_upsert_after_schema_evolution(catalog: Catalog) -> None: + identifier = "default.test_upsert_after_schema_evolution" + _drop_table(catalog, identifier) + + schema = Schema( + NestedField(1, "id", LongType(), required=True), + NestedField(2, "name", StringType(), required=False), + identifier_field_ids=[1], + ) + tbl = catalog.create_table(identifier, schema=schema) + + # Initial write with 2 columns + arrow_schema_v0 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("name", pa.string(), nullable=True), + ] + ) + df_v0 = pa.Table.from_pylist([{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}], schema=arrow_schema_v0) + tbl.append(df_v0) + + # Record initial data file before evolution & upsert + initial_files = [f.file.file_path for f in tbl.scan().plan_files()] + assert len(initial_files) == 1 + + # Schema evolution: add column 'city' + with tbl.update_schema() as update: + update.add_column("city", StringType()) + + # Upsert with 3 columns: update id=1 with new city, insert id=3 + arrow_schema_v1 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("name", pa.string(), nullable=True), + pa.field("city", pa.string(), nullable=True), + ] + ) + df_v1 = pa.Table.from_pylist( + [ + {"id": 1, "name": "Alice", "city": "NYC"}, + {"id": 3, "name": "Charlie", "city": "LA"}, + ], + schema=arrow_schema_v1, + ) + result = tbl.upsert(df_v1) + assert result.rows_updated == 1 + assert result.rows_inserted == 1 + + # Verify that the old V0 data file was replaced (Copy-on-Write overwrite) + current_files = [f.file.file_path for f in tbl.scan().plan_files()] + assert initial_files[0] not in current_files + + # Verify snapshot operations include both OVERWRITE and APPEND + operations = [s.summary.operation for s in tbl.snapshots() if s.summary is not None] + assert Operation.OVERWRITE in operations + assert Operation.APPEND in operations + + # Verify scanned table rows contain updated evolved values + scanned_rows = tbl.scan().to_arrow().to_pylist() + assert sorted(scanned_rows, key=lambda x: x["id"]) == [ + {"id": 1, "name": "Alice", "city": "NYC"}, + {"id": 2, "name": "Bob", "city": None}, + {"id": 3, "name": "Charlie", "city": "LA"}, + ] + + +def test_upsert_after_schema_evolution_union_by_name(catalog: Catalog) -> None: + identifier = "default.test_upsert_after_schema_evolution_union_by_name" + _drop_table(catalog, identifier) + + arrow_schema_v0 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("name", pa.string(), nullable=True), + pa.field("age", pa.int32(), nullable=True), + ] + ) + schema = Schema( + NestedField(1, "id", LongType(), required=True), + NestedField(2, "name", StringType(), required=False), + NestedField(3, "age", IntegerType(), required=False), + identifier_field_ids=[1], + ) + tbl = catalog.create_table(identifier, schema=schema) + + df_v0 = pa.Table.from_pylist( + [ + {"id": 1, "name": "Alice", "age": 30}, + {"id": 2, "name": "Bob", "age": 25}, + ], + schema=arrow_schema_v0, + ) + tbl.append(df_v0) + + # Schema evolution via union_by_name with a new column 'city' + arrow_schema_v1 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("name", pa.string(), nullable=True), + pa.field("age", pa.int32(), nullable=True), + pa.field("city", pa.string(), nullable=True), + ] + ) + with tbl.update_schema() as update: + update.union_by_name(arrow_schema_v1) + + # Upsert with 4 columns: update id=1 (change age & city), unchanged id=2, insert id=3 + df_v1 = pa.Table.from_pylist( + [ + {"id": 1, "name": "Alice", "age": 31, "city": "Taipei"}, + {"id": 2, "name": "Bob", "age": 25, "city": None}, + {"id": 3, "name": "Charlie", "age": 40, "city": "Tokyo"}, + ], + schema=arrow_schema_v1, + ) + result = tbl.upsert(df_v1) + assert result.rows_updated == 1 + assert result.rows_inserted == 1 + + scanned_rows = tbl.scan().to_arrow().to_pylist() + assert sorted(scanned_rows, key=lambda x: x["id"]) == [ + {"id": 1, "name": "Alice", "age": 31, "city": "Taipei"}, + {"id": 2, "name": "Bob", "age": 25, "city": None}, + {"id": 3, "name": "Charlie", "age": 40, "city": "Tokyo"}, + ] + + +def test_upsert_after_multiple_schema_evolutions_with_composite_keys(catalog: Catalog) -> None: + identifier = "default.test_upsert_after_multiple_schema_evolutions_with_composite_keys" + _drop_table(catalog, identifier) + + # Step 1: Create table with V0 schema (id, dept, name) + schema = Schema( + NestedField(1, "id", LongType(), required=True), + NestedField(2, "dept", StringType(), required=True), + NestedField(3, "name", StringType(), required=False), + identifier_field_ids=[1, 2], + ) + tbl = catalog.create_table(identifier, schema=schema) + + arrow_schema_v0 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("dept", pa.string(), nullable=False), + pa.field("name", pa.string(), nullable=True), + ] + ) + tbl.append(pa.Table.from_pylist([{"id": 1, "dept": "ENG", "name": "Alice"}], schema=arrow_schema_v0)) + + # Step 2: Evolve to V1 by adding 'salary' + with tbl.update_schema() as update: + update.add_column("salary", LongType()) + + arrow_schema_v1 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("dept", pa.string(), nullable=False), + pa.field("name", pa.string(), nullable=True), + pa.field("salary", pa.int64(), nullable=True), + ] + ) + tbl.append(pa.Table.from_pylist([{"id": 2, "dept": "HR", "name": "Bob", "salary": 50000}], schema=arrow_schema_v1)) + + # Step 3: Evolve to V2 by adding 'city' + with tbl.update_schema() as update: + update.add_column("city", StringType()) + + arrow_schema_v2 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("dept", pa.string(), nullable=False), + pa.field("name", pa.string(), nullable=True), + pa.field("salary", pa.int64(), nullable=True), + pa.field("city", pa.string(), nullable=True), + ] + ) + + # Step 4: Upsert spanning V0, V1, and V2 rows + df_v2 = pa.Table.from_pylist( + [ + {"id": 1, "dept": "ENG", "name": "Alice", "salary": 80000, "city": "Taipei"}, # Update V0 row (salary + city added) + {"id": 2, "dept": "HR", "name": "Bob", "salary": 50000, "city": "London"}, # Update V1 row (city added) + {"id": 3, "dept": "MKT", "name": "Charlie", "salary": 60000, "city": "Tokyo"}, # Insert new V2 row + ], + schema=arrow_schema_v2, + ) + result = tbl.upsert(df_v2) + assert result.rows_updated == 2 + assert result.rows_inserted == 1 + + scanned_rows = tbl.scan().to_arrow().to_pylist() + assert sorted(scanned_rows, key=lambda x: (x["id"], x["dept"])) == [ + {"id": 1, "dept": "ENG", "name": "Alice", "salary": 80000, "city": "Taipei"}, + {"id": 2, "dept": "HR", "name": "Bob", "salary": 50000, "city": "London"}, + {"id": 3, "dept": "MKT", "name": "Charlie", "salary": 60000, "city": "Tokyo"}, + ] + + +def test_upsert_after_schema_evolution_noop_and_nulls(catalog: Catalog) -> None: + identifier = "default.test_upsert_after_schema_evolution_noop_and_nulls" + _drop_table(catalog, identifier) + + schema = Schema( + NestedField(1, "id", LongType(), required=True), + NestedField(2, "name", StringType(), required=False), + identifier_field_ids=[1], + ) + tbl = catalog.create_table(identifier, schema=schema) + + arrow_schema_v0 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("name", pa.string(), nullable=True), + ] + ) + tbl.append(pa.Table.from_pylist([{"id": 1, "name": "Alice"}], schema=arrow_schema_v0)) + + # Evolve schema + with tbl.update_schema() as update: + update.add_column("extra", StringType()) + + arrow_schema_v1 = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("name", pa.string(), nullable=True), + pa.field("extra", pa.string(), nullable=True), + ] + ) + + # Upsert with identical row where new column is None -> should be no-op (0 updated, 0 inserted) + df_noop = pa.Table.from_pylist([{"id": 1, "name": "Alice", "extra": None}], schema=arrow_schema_v1) + result = tbl.upsert(df_noop) + assert result.rows_updated == 0 + assert result.rows_inserted == 0