feat: Support difference_cols for upsert update detection (#3598)

Author: hedger9487

pyiceberg/table/__init__.py

@@ -848,6 +848,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.
 
@@ -862,6 +863,11 @@ 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. Changes to columns outside difference_cols are
+                intentionally ignored for update detection. 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
 
@@ -913,6 +919,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
@@ -952,7 +961,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
@@ -1653,6 +1662,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.
 
@@ -1667,6 +1677,11 @@ 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. Changes to columns outside difference_cols are
+                intentionally ignored for update detection. 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
 
@@ -1701,6 +1716,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,50 @@ 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.
+
+    These columns must be non-key columns used for change detection.
+
+    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 must contain at least one column; 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 non-key columns are compared to detect changes in
+    matched rows, instead of all non-key columns. Changes to columns outside `difference_cols`
+    are intentionally ignored for update detection. 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")