feat: Support difference_cols for upsert update detection

mkdocs/docs/api.md

@@ -578,6 +578,14 @@ assert upd.rows_inserted == 1
 
 PyIceberg will automatically detect which rows need to be updated, inserted or can simply be ignored.
 
+By default, all non-key columns are compared to detect whether a matched row has changed. When comparing all columns is expensive (for example wide tables, or complex types such as structs and lists), you can limit the comparison to a subset of columns with `difference_cols` — for example a single hash column that reflects any change to the row:
+
+```python
+upd = tbl.upsert(df, difference_cols=["row_hash"])
+```
+
+Note that `difference_cols` only limits change *detection*: when a matched row is detected as changed, all of its columns are written, not just the listed ones.
+
 ## Inspecting tables
 
 To explore the table metadata, tables can be inspected.

pyiceberg/table/__init__.py

@@ -806,6 +806,7 @@ def upsert(
         case_sensitive: bool = True,
         branch: str | None = MAIN_BRANCH,
         snapshot_properties: dict[str, str] = EMPTY_DICT,
+        difference_cols: list[str] | None = None,
     ) -> UpsertResult:
         """Shorthand API for performing an upsert to an iceberg table.
 
@@ -820,6 +821,10 @@ def upsert(
             case_sensitive: Bool indicating if the match should be case-sensitive
             branch: Branch Reference to run the upsert operation
             snapshot_properties: Custom properties to be added to the snapshot summary
+            difference_cols: Subset of non-key columns to compare when detecting changed rows
+                (e.g. a hash column that reflects any change to the row). This only limits change
+                *detection*: when a matched row is detected as changed, all of its columns are
+                written, not just the listed ones. If not provided, all non-key columns are compared.
 
             To learn more about the identifier-field-ids: https://iceberg.apache.org/spec/#identifier-field-ids
 
@@ -871,6 +876,9 @@ def upsert(
         if upsert_util.has_duplicate_rows(df, join_cols):
             raise ValueError("Duplicate rows found in source dataset based on the key columns. No upsert executed")
 
+        # Fail fast on invalid difference_cols instead of erroring on the first matched batch
+        upsert_util.validate_difference_cols(df.column_names, join_cols, difference_cols)
+
         from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible
 
         downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
@@ -910,7 +918,7 @@ def upsert(
                 # values have actually changed. We don't want to do just a blanket overwrite for matched
                 # rows if the actual non-key column data hasn't changed.
                 # this extra step avoids unnecessary IO and writes
-                rows_to_update = upsert_util.get_rows_to_update(df, rows, join_cols)
+                rows_to_update = upsert_util.get_rows_to_update(df, rows, join_cols, difference_cols)
 
                 if len(rows_to_update) > 0:
                     # build the match predicate filter
@@ -1482,6 +1490,7 @@ def upsert(
         case_sensitive: bool = True,
         branch: str | None = MAIN_BRANCH,
         snapshot_properties: dict[str, str] = EMPTY_DICT,
+        difference_cols: list[str] | None = None,
     ) -> UpsertResult:
         """Shorthand API for performing an upsert to an iceberg table.
 
@@ -1496,6 +1505,10 @@ def upsert(
             case_sensitive: Bool indicating if the match should be case-sensitive
             branch: Branch Reference to run the upsert operation
             snapshot_properties: Custom properties to be added to the snapshot summary
+            difference_cols: Subset of non-key columns to compare when detecting changed rows
+                (e.g. a hash column that reflects any change to the row). This only limits change
+                *detection*: when a matched row is detected as changed, all of its columns are
+                written, not just the listed ones. If not provided, all non-key columns are compared.
 
             To learn more about the identifier-field-ids: https://iceberg.apache.org/spec/#identifier-field-ids
 
@@ -1530,6 +1543,7 @@ def upsert(
                 case_sensitive=case_sensitive,
                 branch=branch,
                 snapshot_properties=snapshot_properties,
+                difference_cols=difference_cols,
             )
 
     def append(

pyiceberg/table/upsert_util.py

@@ -53,17 +53,47 @@ def has_duplicate_rows(df: pyarrow_table, join_cols: list[str]) -> bool:
     return len(df.select(join_cols).group_by(join_cols).aggregate([([], "count_all")]).filter(pc.field("count_all") > 1)) > 0
 
 
-def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols: list[str]) -> pa.Table:
+def validate_difference_cols(column_names: list[str], join_cols: list[str], difference_cols: list[str] | None) -> None:
+    """Validate the columns used to detect changes in matched rows.
+
+    Raises:
+        ValueError: If `difference_cols` is empty, contains columns that are not present
+            in `column_names`, or overlaps with `join_cols`.
+    """
+    if difference_cols is None:
+        return
+
+    difference_cols_set = set(difference_cols)
+
+    if not difference_cols_set:
+        raise ValueError("difference_cols cannot be empty, use None to compare all non-key columns")
+
+    if unknown_cols := difference_cols_set - set(column_names):
+        raise ValueError(f"Columns in difference_cols could not be found in the source table: {sorted(unknown_cols)}")
+
+    if key_cols := difference_cols_set & set(join_cols):
+        raise ValueError(f"Columns in difference_cols cannot be join columns: {sorted(key_cols)}")
+
+
+def get_rows_to_update(
+    source_table: pa.Table, target_table: pa.Table, join_cols: list[str], difference_cols: list[str] | None = None
+) -> pa.Table:
     """
     Return a table with rows that need to be updated in the target table based on the join columns.
 
     The table is joined on the identifier columns, and then checked if there are any updated rows.
     Those are selected and everything is renamed correctly.
+
+    When `difference_cols` is provided, only those columns are compared to detect changes in
+    matched rows, instead of all non-key columns. This only affects change *detection*: rows
+    that are detected as changed are still returned with all of their columns.
     """
     all_columns = set(source_table.column_names)
     join_cols_set = set(join_cols)
 
-    non_key_cols = list(all_columns - join_cols_set)
+    validate_difference_cols(source_table.column_names, join_cols, difference_cols)
+
+    non_key_cols = list(all_columns - join_cols_set) if difference_cols is None else difference_cols
 
     if has_duplicate_rows(target_table, join_cols):
         raise ValueError("Target table has duplicate rows, aborting upsert")

tests/table/test_upsert.py

@@ -29,7 +29,7 @@
 from pyiceberg.schema import Schema
 from pyiceberg.table import Table, UpsertResult
 from pyiceberg.table.snapshots import Operation
-from pyiceberg.table.upsert_util import create_match_filter
+from pyiceberg.table.upsert_util import create_match_filter, get_rows_to_update
 from pyiceberg.types import IntegerType, NestedField, StringType, StructType
 from tests.catalog.test_base import InMemoryCatalog
 
@@ -888,3 +888,112 @@ 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_get_rows_to_update_with_difference_cols() -> None:
+    """
+    Change detection is limited to difference_cols, but detected rows are returned with all columns.
+    """
+    schema = pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string()), pa.field("meta", pa.string())])
+    target = pa.Table.from_pylist(
+        [
+            {"id": 1, "val": "a", "meta": "m1"},
+            {"id": 2, "val": "b", "meta": "m2"},
+        ],
+        schema=schema,
+    )
+    source = pa.Table.from_pylist(
+        [
+            {"id": 1, "val": "a", "meta": "changed"},  # differs only in a column outside difference_cols
+            {"id": 2, "val": "B", "meta": "m2"},  # differs in a difference_cols column
+        ],
+        schema=schema,
+    )
+
+    # Without difference_cols, both rows are detected as changed
+    assert len(get_rows_to_update(source, target, ["id"])) == 2
+
+    # With difference_cols, only the row with a change in "val" is detected,
+    # and it is returned with all of its columns
+    rows = get_rows_to_update(source, target, ["id"], difference_cols=["val"])
+    assert rows.to_pylist() == [{"id": 2, "val": "B", "meta": "m2"}]
+
+
+def test_get_rows_to_update_difference_cols_validation() -> None:
+    schema = pa.schema([pa.field("id", pa.int32()), pa.field("val", pa.string())])
+    table = pa.Table.from_pylist([{"id": 1, "val": "a"}], schema=schema)
+
+    with pytest.raises(ValueError, match="could not be found in the source table"):
+        get_rows_to_update(table, table, ["id"], difference_cols=["nonexistent"])
+
+    with pytest.raises(ValueError, match="cannot be join columns"):
+        get_rows_to_update(table, table, ["id"], difference_cols=["id"])
+
+    with pytest.raises(ValueError, match="cannot be empty"):
+        get_rows_to_update(table, table, ["id"], difference_cols=[])
+
+
+def test_upsert_with_difference_cols(catalog: Catalog) -> None:
+    """
+    Upsert with difference_cols skips matched rows whose changes are outside the listed columns,
+    while updated rows are written with all of their columns.
+    """
+    identifier = "default.test_upsert_with_difference_cols"
+    _drop_table(catalog, identifier)
+
+    arrow_schema = pa.schema(
+        [
+            pa.field("city", pa.string(), nullable=False),
+            pa.field("population", pa.int32(), nullable=False),
+            pa.field("notes", pa.string(), nullable=False),
+        ]
+    )
+
+    tbl = catalog.create_table(identifier, arrow_schema)
+    tbl.append(
+        pa.Table.from_pylist(
+            [
+                {"city": "Amsterdam", "population": 921402, "notes": "old"},
+                {"city": "San Francisco", "population": 808988, "notes": "old"},
+            ],
+            schema=arrow_schema,
+        )
+    )
+
+    source_df = pa.Table.from_pylist(
+        [
+            {"city": "Amsterdam", "population": 921402, "notes": "new"},  # change outside difference_cols -> skipped
+            {"city": "San Francisco", "population": 810000, "notes": "new"},  # change in difference_cols -> updated
+            {"city": "Drachten", "population": 45019, "notes": "new"},  # unmatched -> inserted
+        ],
+        schema=arrow_schema,
+    )
+
+    res = tbl.upsert(source_df, join_cols=["city"], difference_cols=["population"])
+
+    assert_upsert_result(res, expected_updated=1, expected_inserted=1)
+
+    result = {row["city"]: row for row in tbl.scan().to_arrow().to_pylist()}
+    # The skipped row is untouched, including the column that differed
+    assert result["Amsterdam"] == {"city": "Amsterdam", "population": 921402, "notes": "old"}
+    # The updated row is written with all columns, not only the difference_cols
+    assert result["San Francisco"] == {"city": "San Francisco", "population": 810000, "notes": "new"}
+    assert result["Drachten"] == {"city": "Drachten", "population": 45019, "notes": "new"}
+
+
+def test_upsert_with_invalid_difference_cols(catalog: Catalog) -> None:
+    identifier = "default.test_upsert_with_invalid_difference_cols"
+    _drop_table(catalog, identifier)
+
+    arrow_schema = pa.schema(
+        [
+            pa.field("city", pa.string(), nullable=False),
+            pa.field("population", pa.int32(), nullable=False),
+        ]
+    )
+
+    tbl = catalog.create_table(identifier, arrow_schema)
+    df = pa.Table.from_pylist([{"city": "Amsterdam", "population": 921402}], schema=arrow_schema)
+
+    with pytest.raises(ValueError, match="could not be found in the source table"):
+        tbl.upsert(df, join_cols=["city"], difference_cols=["nonexistent"])