From 0a9ddafa8399c7e6667d305c2059d3d1860cb8c8 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Wed, 19 Aug 2026 12:54:41 -0700 Subject: [PATCH 1/3] IO: Absorb upsert PyArrow logic into io/pyarrow.py Move PyArrow-specific table operations (joins, group_by, duplicate detection, row comparison) from table/upsert_util.py into pyiceberg/io/pyarrow.py. The upsert_util module now delegates to helpers in io/pyarrow rather than importing pyarrow directly. This consolidates PyArrow logic behind the io/pyarrow module boundary, which is a precondition for the decomposition proposed in #3737. Fixes #3812 (PR A) --- pyiceberg/io/pyarrow.py | 82 ++++++++++++++++++++++++++++++ pyiceberg/table/upsert_util.py | 92 +++++++--------------------------- 2 files changed, 100 insertions(+), 74 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..14a889791c 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -3137,3 +3137,85 @@ 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_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/upsert_util.py b/pyiceberg/table/upsert_util.py index 6f32826eb0..d104178a0f 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -16,10 +16,7 @@ # under the License. import functools import operator - -import pyarrow as pa -from pyarrow import Table as pyarrow_table -from pyarrow import compute as pc +from typing import TYPE_CHECKING from pyiceberg.expressions import ( AlwaysFalse, @@ -28,10 +25,19 @@ In, Or, ) +from pyiceberg.io.pyarrow import ( + _upsert_get_rows_to_update, + _upsert_has_duplicate_rows, + _upsert_unique_keys, +) +if TYPE_CHECKING: + import pyarrow as pa -def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpression: - unique_keys = df.select(join_cols).group_by(join_cols).aggregate([]) + +def 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()) @@ -48,77 +54,15 @@ def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpre return Or(*filters) -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 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) -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. +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 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() - - # 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 - - # 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)))) - - # 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 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() + return _upsert_get_rows_to_update(source_table, target_table, join_cols) From 73392efbcda01f3b05a66e9d7b83b6b3b1fe170d Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Fri, 21 Aug 2026 11:12:08 -0700 Subject: [PATCH 2/3] Drop leading underscore from upsert helpers in io/pyarrow.py --- pyiceberg/io/pyarrow.py | 8 ++++---- pyiceberg/table/upsert_util.py | 12 ++++++------ 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 14a889791c..5992836ee0 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -3139,7 +3139,7 @@ def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Ar return pc.struct_field(field_array, path_parts[1:]) -def _upsert_unique_keys(df: pa.Table, join_cols: list[str]) -> pa.Table: +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. @@ -3147,12 +3147,12 @@ def _upsert_unique_keys(df: pa.Table, join_cols: list[str]) -> pa.Table: return df.select(join_cols).group_by(join_cols).aggregate([]) -def _upsert_has_duplicate_rows(df: pa.Table, join_cols: list[str]) -> bool: +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: +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 @@ -3168,7 +3168,7 @@ def _upsert_get_rows_to_update(source_table: pa.Table, target_table: pa.Table, j non_key_cols = list(all_columns - join_cols_set) - if _upsert_has_duplicate_rows(target_table, join_cols): + if upsert_has_duplicate_rows(target_table, join_cols): raise ValueError("Target table has duplicate rows, aborting upsert") if len(target_table) == 0: diff --git a/pyiceberg/table/upsert_util.py b/pyiceberg/table/upsert_util.py index d104178a0f..908024d8ba 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -26,9 +26,9 @@ Or, ) from pyiceberg.io.pyarrow import ( - _upsert_get_rows_to_update, - _upsert_has_duplicate_rows, - _upsert_unique_keys, + upsert_get_rows_to_update, + upsert_has_duplicate_rows, + upsert_unique_keys, ) if TYPE_CHECKING: @@ -37,7 +37,7 @@ def 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) + unique_keys = upsert_unique_keys(df, join_cols) if len(join_cols) == 1: return In(join_cols[0], unique_keys[0].to_pylist()) @@ -56,7 +56,7 @@ def create_match_filter(df: "pa.Table", join_cols: list[str]) -> BooleanExpressi 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) + return upsert_has_duplicate_rows(df, join_cols) def get_rows_to_update(source_table: "pa.Table", target_table: "pa.Table", join_cols: list[str]) -> "pa.Table": @@ -65,4 +65,4 @@ def get_rows_to_update(source_table: "pa.Table", target_table: "pa.Table", join_ 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. """ - return _upsert_get_rows_to_update(source_table, target_table, join_cols) + return upsert_get_rows_to_update(source_table, target_table, join_cols) From c3754e6ac34d200bc2845340e65255f7de74c3c2 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Wed, 9 Sep 2026 15:40:41 -0700 Subject: [PATCH 3/3] IO: Dissolve upsert_util.py into io/pyarrow.py with deprecation shim --- pyiceberg/io/pyarrow.py | 32 ++++++++++++++++++- pyiceberg/table/__init__.py | 23 ++++++++------ pyiceberg/table/upsert_util.py | 58 +++++++++++++++------------------- tests/table/test_upsert.py | 5 ++- 4 files changed, 72 insertions(+), 46 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 5992836ee0..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, @@ -3147,6 +3158,25 @@ def upsert_unique_keys(df: pa.Table, join_cols: list[str]) -> pa.Table: 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 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 908024d8ba..dd30de00d5 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -14,55 +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 + +"""Deprecated: upsert helpers have moved to pyiceberg.io.pyarrow. + +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. +""" + +from __future__ import annotations + from typing import TYPE_CHECKING -from pyiceberg.expressions import ( - AlwaysFalse, - BooleanExpression, - EqualTo, - In, - Or, -) +from pyiceberg.expressions import BooleanExpression from pyiceberg.io.pyarrow import ( + upsert_create_match_filter, upsert_get_rows_to_update, upsert_has_duplicate_rows, - upsert_unique_keys, ) +from pyiceberg.utils.deprecated import deprecated if TYPE_CHECKING: import pyarrow as pa +_DEPRECATION_IN = "0.13.0" +_REMOVAL_IN = "0.14.0" +_HELP = "Use the equivalent function from pyiceberg.io.pyarrow instead" -def 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) +@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) -def has_duplicate_rows(df: "pa.Table", join_cols: list[str]) -> bool: +@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) -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 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. - """ +@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)),