diff --git a/pyiceberg/table/upsert_util.py b/pyiceberg/table/upsert_util.py index 6f32826eb0..1a5692fa69 100644 --- a/pyiceberg/table/upsert_util.py +++ b/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() diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..4291c18a55 100644 --- a/tests/table/test_upsert.py +++ b/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}]}