Core: Compare upsert rows column by column instead of cell by cell

pyiceberg/table/upsert_util.py

@@ -53,6 +53,25 @@ 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 _values_differ(
+    source_values: pa.Array | pa.ChunkedArray, target_values: pa.Array | pa.ChunkedArray
+) -> pa.Array | pa.ChunkedArray:
+    """Return a mask of the positions where the values are not equal, treating two nulls as equal."""
+    try:
+        equal = pc.fill_null(pc.equal(source_values, target_values), False)
+    except (pa.ArrowNotImplementedError, pa.ArrowInvalid, pa.ArrowTypeError):
+        # PyArrow has no comparison kernel for nested types
+        source_list = source_values.to_pylist()
+        target_list = target_values.to_pylist()
+        return pa.array(
+            [source_value != target_value for source_value, target_value in zip(source_list, target_list, strict=True)],
+            type=pa.bool_(),
+        )
+
+    both_null = pc.and_(pc.is_null(source_values), pc.is_null(target_values))
+    return pc.invert(pc.or_(equal, both_null))
+
+
 def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols: list[str]) -> pa.Table:
     """
     Return a table with rows that need to be updated in the target table based on the join columns.
@@ -72,7 +91,7 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols
         # When the target table is empty, there is nothing to update :)
         return source_table.schema.empty_table()
 
-    # We need to compare non_key_cols in Python as PyArrow
+    # We need to compare non_key_cols without a join as PyArrow
     # 1. Cannot do a join when non-join columns have complex types
     # 2. Cannot compare columns with complex types
     # See: https://github.com/apache/arrow/issues/35785
@@ -100,25 +119,20 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols
     # 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")
 
-    # Step 4: Compare all rows using Python
-    to_update_indices = []
-    for source_idx, target_idx in zip(
-        matching_indices[SOURCE_INDEX_COLUMN_NAME].to_pylist(),
-        matching_indices[TARGET_INDEX_COLUMN_NAME].to_pylist(),
-        strict=True,
-    ):
-        source_row = source_table.slice(source_idx, 1)
-        target_row = target_table.slice(target_idx, 1)
-
-        for key in non_key_cols:
-            source_val = source_row.column(key)[0].as_py()
-            target_val = target_row.column(key)[0].as_py()
-            if source_val != target_val:
-                to_update_indices.append(source_idx)
-                break
+    # Step 4: Compare the matched rows column by column
+    source_indices = matching_indices[SOURCE_INDEX_COLUMN_NAME]
+    target_indices = matching_indices[TARGET_INDEX_COLUMN_NAME]
+
+    changed: pa.Array | pa.ChunkedArray | None = None
+    for key in non_key_cols:
+        column_changed = _values_differ(
+            source_table.column(key).take(source_indices),
+            target_table.column(key).take(target_indices),
+        )
+        changed = column_changed if changed is None else pc.or_(changed, column_changed)
 
     # Step 5: Take rows from source table using the indices and cast to target schema
-    if to_update_indices:
-        return source_table.take(to_update_indices)
+    if changed is not None and pc.any(changed).as_py():
+        return source_table.take(pc.filter(source_indices, changed))
     else:
         return source_table.schema.empty_table()

tests/table/test_upsert.py

@@ -31,7 +31,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.transforms import DayTransform
 from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType
 from tests.catalog.test_base import InMemoryCatalog
@@ -927,3 +927,31 @@ 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_returns_changed_rows() -> None:
+    source = pa.table({"pk": [1, 2, 3], "value": [10, 20, 30]})
+    target = pa.table({"pk": [1, 2, 3], "value": [10, 21, 30]})
+
+    assert get_rows_to_update(source, target, ["pk"]).to_pydict() == {"pk": [2], "value": [20]}
+
+
+def test_get_rows_to_update_returns_a_row_once() -> None:
+    source = pa.table({"pk": [1], "a": [1], "b": [2]})
+    target = pa.table({"pk": [1], "a": [9], "b": [9]})
+
+    assert get_rows_to_update(source, target, ["pk"]).to_pydict() == {"pk": [1], "a": [1], "b": [2]}
+
+
+def test_get_rows_to_update_treats_nulls_as_equal() -> None:
+    source = pa.table({"pk": [1, 2], "value": pa.array([None, 1], type=pa.int64())})
+    target = pa.table({"pk": [1, 2], "value": pa.array([None, 2], type=pa.int64())})
+
+    assert get_rows_to_update(source, target, ["pk"]).to_pydict() == {"pk": [2], "value": [1]}
+
+
+def test_get_rows_to_update_compares_nested_values() -> None:
+    source = pa.table({"pk": [1, 2], "value": pa.array([{"a": 1}, {"a": 2}])})
+    target = pa.table({"pk": [1, 2], "value": pa.array([{"a": 1}, {"a": 3}])})
+
+    assert get_rows_to_update(source, target, ["pk"]).to_pydict() == {"pk": [2], "value": [{"a": 2}]}