Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion pyiceberg/table/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand All @@ -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

Expand Down Expand Up @@ -1701,6 +1716,7 @@ def upsert(
case_sensitive=case_sensitive,
branch=branch,
snapshot_properties=snapshot_properties,
difference_cols=difference_cols,
)

def append(
Expand Down
37 changes: 35 additions & 2 deletions pyiceberg/table/upsert_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Loading