IO: Dissolve upsert_util.py into io/pyarrow.py with deprecation shim

Author: qzyu999

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

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:

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)

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)),