diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..587bf81cfc 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -68,7 +68,18 @@ from pyiceberg.conversions import to_bytes from pyiceberg.exceptions import ResolveError -from pyiceberg.expressions import AlwaysTrue, BooleanExpression, BoundIsNaN, BoundIsNull, BoundTerm, Not, Or +from pyiceberg.expressions import ( + AlwaysFalse, + AlwaysTrue, + BooleanExpression, + BoundIsNaN, + BoundIsNull, + BoundTerm, + EqualTo, + In, + Not, + Or, +) from pyiceberg.expressions.literals import Literal from pyiceberg.expressions.visitors import ( BoundBooleanExpressionVisitor, @@ -3137,3 +3148,104 @@ def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Ar field_array = arrow_table[path_parts[0]] # Navigate into the struct using the remaining path parts return pc.struct_field(field_array, path_parts[1:]) + + +def upsert_unique_keys(df: pa.Table, join_cols: list[str]) -> pa.Table: + """Extract unique key combinations from a table. + + Returns a table containing one row per distinct combination of join_cols. + """ + return df.select(join_cols).group_by(join_cols).aggregate([]) + + +def upsert_create_match_filter(df: pa.Table, join_cols: list[str]) -> BooleanExpression: + """Build an Iceberg filter expression matching the unique keys in df.""" + unique_keys = upsert_unique_keys(df, join_cols) + + if len(join_cols) == 1: + return In(join_cols[0], unique_keys[0].to_pylist()) + else: + filters = [ + functools.reduce(operator.and_, [EqualTo(col, row[col]) for col in join_cols]) for row in unique_keys.to_pylist() + ] + + if len(filters) == 0: + return AlwaysFalse() + elif len(filters) == 1: + return filters[0] + else: + return Or(*filters) + + +def upsert_has_duplicate_rows(df: pa.Table, join_cols: list[str]) -> bool: + """Check for duplicate rows in a PyArrow table based on the join columns.""" + return len(df.select(join_cols).group_by(join_cols).aggregate([([], "count_all")]).filter(pc.field("count_all") > 1)) > 0 + + +def upsert_get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols: list[str]) -> pa.Table: + """Return rows from source_table whose non-key columns differ from target_table. + + Performs an inner join on join_cols, then compares non-key column values + row-by-row. Returns the subset of source rows that have at least one + changed non-key column. If target_table is empty, returns an empty table. + + Raises: + ValueError: If target_table has duplicate rows on join_cols. + ValueError: If join_cols use reserved index column names. + """ + all_columns = set(source_table.column_names) + join_cols_set = set(join_cols) + + non_key_cols = list(all_columns - join_cols_set) + + if upsert_has_duplicate_rows(target_table, join_cols): + raise ValueError("Target table has duplicate rows, aborting upsert") + + if len(target_table) == 0: + return source_table.schema.empty_table() + + # We need to compare non_key_cols in Python 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 + SOURCE_INDEX_COLUMN_NAME = "__source_index" + TARGET_INDEX_COLUMN_NAME = "__target_index" + + if SOURCE_INDEX_COLUMN_NAME in join_cols or TARGET_INDEX_COLUMN_NAME in join_cols: + raise ValueError( + f"{SOURCE_INDEX_COLUMN_NAME} and {TARGET_INDEX_COLUMN_NAME} are reserved for joining " + f"DataFrames, and cannot be used as column names" + ) from None + + # Cast to target table schema so types align for the join. + # See: https://github.com/apache/arrow/issues/37542 + source_index = ( + source_table.cast(target_table.schema) + .select(join_cols_set) + .append_column(SOURCE_INDEX_COLUMN_NAME, pa.array(range(len(source_table)))) + ) + + target_index = target_table.select(join_cols_set).append_column(TARGET_INDEX_COLUMN_NAME, pa.array(range(len(target_table)))) + + matching_indices = source_index.join(target_index, keys=list(join_cols_set), join_type="inner") + + 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 + + if to_update_indices: + return source_table.take(to_update_indices) + else: + return source_table.schema.empty_table() diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index fca718f5ec..9536be34ff 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -892,8 +892,12 @@ def upsert( except ModuleNotFoundError as e: raise ModuleNotFoundError("For writes PyArrow needs to be installed") from e - from pyiceberg.io.pyarrow import expression_to_pyarrow - from pyiceberg.table import upsert_util + from pyiceberg.io.pyarrow import ( + expression_to_pyarrow, + upsert_create_match_filter, + upsert_get_rows_to_update, + upsert_has_duplicate_rows, + ) if join_cols is None: join_cols = [] @@ -910,7 +914,7 @@ def upsert( if not when_matched_update_all and not when_not_matched_insert_all: raise ValueError("no upsert options selected...exiting") - if upsert_util.has_duplicate_rows(df, join_cols): + if upsert_has_duplicate_rows(df, join_cols): raise ValueError("Duplicate rows found in source dataset based on the key columns. No upsert executed") from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible @@ -924,7 +928,7 @@ def upsert( ) # get list of rows that exist so we don't have to load the entire target table - matched_predicate = upsert_util.create_match_filter(df, join_cols) + matched_predicate = upsert_create_match_filter(df, join_cols) # We must use Transaction.table_metadata for the scan. This includes all uncommitted - but relevant - changes. @@ -952,17 +956,17 @@ 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_get_rows_to_update(df, rows, join_cols) if len(rows_to_update) > 0: # build the match predicate filter - overwrite_mask_predicate = upsert_util.create_match_filter(rows_to_update, join_cols) + overwrite_mask_predicate = upsert_create_match_filter(rows_to_update, join_cols) batches_to_overwrite.append(rows_to_update) overwrite_predicates.append(overwrite_mask_predicate) if when_not_matched_insert_all: - expr_match = upsert_util.create_match_filter(rows, join_cols) + expr_match = upsert_create_match_filter(rows, join_cols) expr_match_bound = bind(self.table_metadata.schema(), expr_match, case_sensitive=case_sensitive) expr_match_arrow = expression_to_pyarrow(expr_match_bound) @@ -2663,8 +2667,9 @@ def plan_files(self) -> Iterable[FileScanTask]: options=self.options, ).plan_files( manifests=manifests, - manifest_entry_filter=lambda manifest_entry: manifest_entry.snapshot_id in append_snapshot_ids - and manifest_entry.status == ManifestEntryStatus.ADDED, + manifest_entry_filter=lambda manifest_entry: ( + manifest_entry.snapshot_id in append_snapshot_ids and manifest_entry.status == ManifestEntryStatus.ADDED + ), ) def to_arrow(self) -> pa.Table: diff --git a/pyiceberg/table/upsert_util.py b/pyiceberg/table/upsert_util.py index 6f32826eb0..dd30de00d5 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -14,111 +14,47 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import functools -import operator -import pyarrow as pa -from pyarrow import Table as pyarrow_table -from pyarrow import compute as pc +"""Deprecated: upsert helpers have moved to pyiceberg.io.pyarrow. -from pyiceberg.expressions import ( - AlwaysFalse, - BooleanExpression, - EqualTo, - In, - Or, -) - - -def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpression: - unique_keys = df.select(join_cols).group_by(join_cols).aggregate([]) - - if len(join_cols) == 1: - return In(join_cols[0], unique_keys[0].to_pylist()) - else: - filters = [ - functools.reduce(operator.and_, [EqualTo(col, row[col]) for col in join_cols]) for row in unique_keys.to_pylist() - ] +All functions in this module are re-exported from ``pyiceberg.io.pyarrow`` +and will emit a ``DeprecationWarning`` when called. Import directly from +``pyiceberg.io.pyarrow`` instead. +""" - if len(filters) == 0: - return AlwaysFalse() - elif len(filters) == 1: - return filters[0] - else: - return Or(*filters) +from __future__ import annotations +from typing import TYPE_CHECKING -def has_duplicate_rows(df: pyarrow_table, join_cols: list[str]) -> bool: - """Check for duplicate rows in a PyArrow table based on the join columns.""" - 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: - """ - 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. - """ - all_columns = set(source_table.column_names) - join_cols_set = set(join_cols) - - non_key_cols = list(all_columns - join_cols_set) - - if has_duplicate_rows(target_table, join_cols): - raise ValueError("Target table has duplicate rows, aborting upsert") - - if len(target_table) == 0: - # When the target table is empty, there is nothing to update :) - return source_table.schema.empty_table() +from pyiceberg.expressions import BooleanExpression +from pyiceberg.io.pyarrow import ( + upsert_create_match_filter, + upsert_get_rows_to_update, + upsert_has_duplicate_rows, +) +from pyiceberg.utils.deprecated import deprecated - # We need to compare non_key_cols in Python 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 - SOURCE_INDEX_COLUMN_NAME = "__source_index" - TARGET_INDEX_COLUMN_NAME = "__target_index" +if TYPE_CHECKING: + import pyarrow as pa - if SOURCE_INDEX_COLUMN_NAME in join_cols or TARGET_INDEX_COLUMN_NAME in join_cols: - raise ValueError( - f"{SOURCE_INDEX_COLUMN_NAME} and {TARGET_INDEX_COLUMN_NAME} are reserved for joining " - f"DataFrames, and cannot be used as column names" - ) from None +_DEPRECATION_IN = "0.13.0" +_REMOVAL_IN = "0.14.0" +_HELP = "Use the equivalent function from pyiceberg.io.pyarrow instead" - # Step 1: Prepare source index with join keys and a marker index - # Cast to target table schema, so we can do the join - # See: https://github.com/apache/arrow/issues/37542 - source_index = ( - source_table.cast(target_table.schema) - .select(join_cols_set) - .append_column(SOURCE_INDEX_COLUMN_NAME, pa.array(range(len(source_table)))) - ) - # Step 2: Prepare target index with join keys and a marker - target_index = target_table.select(join_cols_set).append_column(TARGET_INDEX_COLUMN_NAME, pa.array(range(len(target_table)))) +@deprecated(deprecated_in=_DEPRECATION_IN, removed_in=_REMOVAL_IN, help_message=_HELP) +def create_match_filter(df: pa.Table, join_cols: list[str]) -> BooleanExpression: + """Build an Iceberg filter expression matching the unique keys in df.""" + return upsert_create_match_filter(df, 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) +@deprecated(deprecated_in=_DEPRECATION_IN, removed_in=_REMOVAL_IN, help_message=_HELP) +def has_duplicate_rows(df: pa.Table, join_cols: list[str]) -> bool: + """Check for duplicate rows in a table based on the join columns.""" + return upsert_has_duplicate_rows(df, join_cols) - 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 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) - else: - return source_table.schema.empty_table() +@deprecated(deprecated_in=_DEPRECATION_IN, removed_in=_REMOVAL_IN, help_message=_HELP) +def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols: list[str]) -> pa.Table: + """Return rows from source that need to be updated in the target table.""" + return upsert_get_rows_to_update(source_table, target_table, join_cols) diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..f5d548f1d8 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -26,12 +26,11 @@ from pyiceberg.exceptions import NoSuchTableError from pyiceberg.expressions import AlwaysTrue, And, EqualTo, Reference from pyiceberg.expressions.literals import LongLiteral -from pyiceberg.io.pyarrow import schema_to_pyarrow +from pyiceberg.io.pyarrow import schema_to_pyarrow, upsert_create_match_filter from pyiceberg.partitioning import PartitionField, PartitionSpec 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.transforms import DayTransform from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType from tests.catalog.test_base import InMemoryCatalog @@ -439,7 +438,7 @@ def test_create_match_filter_single_condition() -> None: ] schema = pa.schema([pa.field("order_id", pa.int32()), pa.field("order_line_id", pa.int32()), pa.field("extra", pa.string())]) table = pa.Table.from_pylist(data, schema=schema) - expr = create_match_filter(table, ["order_id", "order_line_id"]) + expr = upsert_create_match_filter(table, ["order_id", "order_line_id"]) assert expr == And( EqualTo(term=Reference(name="order_id"), literal=LongLiteral(101)), EqualTo(term=Reference(name="order_line_id"), literal=LongLiteral(1)),