Upsert: Reuse existing expression to detect rows to be inserted (#1662)

Also slight refactor of the tests to bring it more in line with the rest

Author: Fokko

pyiceberg/table/__init__.py

@@ -62,7 +62,7 @@
     manifest_evaluator,
 )
 from pyiceberg.io import FileIO, load_file_io
-from pyiceberg.io.pyarrow import ArrowScan, schema_to_pyarrow
+from pyiceberg.io.pyarrow import ArrowScan, expression_to_pyarrow, schema_to_pyarrow
 from pyiceberg.manifest import (
     POSITIONAL_DELETE_SCHEMA,
     DataFile,
@@ -1101,7 +1101,12 @@ def name_mapping(self) -> Optional[NameMapping]:
         return self.metadata.name_mapping()
 
     def upsert(
-        self, df: pa.Table, join_cols: list[str], when_matched_update_all: bool = True, when_not_matched_insert_all: bool = True
+        self,
+        df: pa.Table,
+        join_cols: list[str],
+        when_matched_update_all: bool = True,
+        when_not_matched_insert_all: bool = True,
+        case_sensitive: bool = True,
     ) -> UpsertResult:
         """Shorthand API for performing an upsert to an iceberg table.
 
@@ -1111,6 +1116,7 @@ def upsert(
             join_cols: The columns to join on. These are essentially analogous to primary keys
             when_matched_update_all: Bool indicating to update rows that are matched but require an update due to a value in a non-key column changing
             when_not_matched_insert_all: Bool indicating new rows to be inserted that do not match any existing rows in the table
+            case_sensitive: Bool indicating if the match should be case-sensitive
 
                 Example Use Cases:
                     Case 1: Both Parameters = True (Full Upsert)
@@ -1144,7 +1150,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_iceberg_table = self.scan(row_filter=matched_predicate).to_arrow()
+        matched_iceberg_table = self.scan(row_filter=matched_predicate, case_sensitive=case_sensitive).to_arrow()
 
         update_row_cnt = 0
         insert_row_cnt = 0
@@ -1164,7 +1170,10 @@ def upsert(
                 tx.overwrite(rows_to_update, overwrite_filter=overwrite_mask_predicate)
 
             if when_not_matched_insert_all:
-                rows_to_insert = upsert_util.get_rows_to_insert(df, matched_iceberg_table, join_cols)
+                expr_match = upsert_util.create_match_filter(matched_iceberg_table, join_cols)
+                expr_match_bound = bind(self.schema(), expr_match, case_sensitive=case_sensitive)
+                expr_match_arrow = expression_to_pyarrow(expr_match_bound)
+                rows_to_insert = df.filter(~expr_match_arrow)
 
                 insert_row_cnt = len(rows_to_insert)
 

pyiceberg/table/upsert_util.py

@@ -92,27 +92,3 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols
     rows_to_update_table = rows_to_update_table.select(list(common_columns))
 
     return rows_to_update_table
-
-
-def get_rows_to_insert(source_table: pa.Table, target_table: pa.Table, join_cols: list[str]) -> pa.Table:
-    source_filter_expr = pc.scalar(True)
-
-    for col in join_cols:
-        target_values = target_table.column(col).to_pylist()
-        expr = pc.field(col).isin(target_values)
-
-        if source_filter_expr is None:
-            source_filter_expr = expr
-        else:
-            source_filter_expr = source_filter_expr & expr
-
-    non_matching_expr = ~source_filter_expr
-
-    source_columns = set(source_table.column_names)
-    target_columns = set(target_table.column_names)
-
-    common_columns = source_columns.intersection(target_columns)
-
-    non_matching_rows = source_table.filter(non_matching_expr).select(common_columns)
-
-    return non_matching_rows

tests/table/test_upsert.py

@@ -14,14 +14,30 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from pathlib import PosixPath
+
 import pytest
 from datafusion import SessionContext
 from pyarrow import Table as pa_table
 
+from pyiceberg.catalog import Catalog
+from pyiceberg.exceptions import NoSuchTableError
 from pyiceberg.table import UpsertResult
 from tests.catalog.test_base import InMemoryCatalog, Table
 
-_TEST_NAMESPACE = "test_ns"
+
+@pytest.fixture
+def catalog(tmp_path: PosixPath) -> InMemoryCatalog:
+    catalog = InMemoryCatalog("test.in_memory.catalog", warehouse=tmp_path.absolute().as_posix())
+    catalog.create_namespace("default")
+    return catalog
+
+
+def _drop_table(catalog: Catalog, identifier: str) -> None:
+    try:
+        catalog.drop_table(identifier)
+    except NoSuchTableError:
+        pass
 
 
 def show_iceberg_table(table: Table, ctx: SessionContext) -> None:
@@ -72,7 +88,7 @@ def gen_source_dataset(start_row: int, end_row: int, composite_key: bool, add_du
 
 
 def gen_target_iceberg_table(
-    start_row: int, end_row: int, composite_key: bool, ctx: SessionContext, catalog: InMemoryCatalog, namespace: str
+    start_row: int, end_row: int, composite_key: bool, ctx: SessionContext, catalog: InMemoryCatalog, identifier: str
 ) -> Table:
     additional_columns = ", t.order_id + 1000 as order_line_id" if composite_key else ""
 
@@ -83,7 +99,7 @@ def gen_target_iceberg_table(
         from t
     """).to_arrow_table()
 
-    table = catalog.create_table(f"{_TEST_NAMESPACE}.target", df.schema)
+    table = catalog.create_table(identifier, df.schema)
 
     table.append(df)
 
@@ -95,13 +111,6 @@ def assert_upsert_result(res: UpsertResult, expected_updated: int, expected_inse
     assert res.rows_inserted == expected_inserted, f"rows inserted should be {expected_inserted}, but got {res.rows_inserted}"
 
 
-@pytest.fixture(scope="session")
-def catalog_conn() -> InMemoryCatalog:
-    catalog = InMemoryCatalog("test")
-    catalog.create_namespace(namespace=_TEST_NAMESPACE)
-    yield catalog
-
-
 @pytest.mark.parametrize(
     "join_cols, src_start_row, src_end_row, target_start_row, target_end_row, when_matched_update_all, when_not_matched_insert_all, expected_updated, expected_inserted",
     [
@@ -112,7 +121,7 @@ def catalog_conn() -> InMemoryCatalog:
     ],
 )
 def test_merge_rows(
-    catalog_conn: InMemoryCatalog,
+    catalog: Catalog,
     join_cols: list[str],
     src_start_row: int,
     src_end_row: int,
@@ -123,12 +132,13 @@ def test_merge_rows(
     expected_updated: int,
     expected_inserted: int,
 ) -> None:
-    ctx = SessionContext()
+    identifier = "default.test_merge_rows"
+    _drop_table(catalog, identifier)
 
-    catalog = catalog_conn
+    ctx = SessionContext()
 
     source_df = gen_source_dataset(src_start_row, src_end_row, False, False, ctx)
-    ice_table = gen_target_iceberg_table(target_start_row, target_end_row, False, ctx, catalog, _TEST_NAMESPACE)
+    ice_table = gen_target_iceberg_table(target_start_row, target_end_row, False, ctx, catalog, identifier)
     res = ice_table.upsert(
         df=source_df,
         join_cols=join_cols,
@@ -138,13 +148,13 @@ def test_merge_rows(
 
     assert_upsert_result(res, expected_updated, expected_inserted)
 
-    catalog.drop_table(f"{_TEST_NAMESPACE}.target")
 
-
-def test_merge_scenario_skip_upd_row(catalog_conn: InMemoryCatalog) -> None:
+def test_merge_scenario_skip_upd_row(catalog: Catalog) -> None:
     """
     tests a single insert and update; skips a row that does not need to be updated
     """
+    identifier = "default.test_merge_scenario_skip_upd_row"
+    _drop_table(catalog, identifier)
 
     ctx = SessionContext()
 
@@ -154,8 +164,7 @@ def test_merge_scenario_skip_upd_row(catalog_conn: InMemoryCatalog) -> None:
         select 2 as order_id, date '2021-01-01' as order_date, 'A' as order_type
     """).to_arrow_table()
 
-    catalog = catalog_conn
-    table = catalog.create_table(f"{_TEST_NAMESPACE}.target", df.schema)
+    table = catalog.create_table(identifier, df.schema)
 
     table.append(df)
 
@@ -174,24 +183,24 @@ def test_merge_scenario_skip_upd_row(catalog_conn: InMemoryCatalog) -> None:
 
     assert_upsert_result(res, expected_updated, expected_inserted)
 
-    catalog.drop_table(f"{_TEST_NAMESPACE}.target")
-
 
-def test_merge_scenario_date_as_key(catalog_conn: InMemoryCatalog) -> None:
+def test_merge_scenario_date_as_key(catalog: Catalog) -> None:
     """
     tests a single insert and update; primary key is a date column
     """
 
     ctx = SessionContext()
 
+    identifier = "default.test_merge_scenario_date_as_key"
+    _drop_table(catalog, identifier)
+
     df = ctx.sql("""
         select date '2021-01-01' as order_date, 'A' as order_type
         union all
         select date '2021-01-02' as order_date, 'A' as order_type
     """).to_arrow_table()
 
-    catalog = catalog_conn
-    table = catalog.create_table(f"{_TEST_NAMESPACE}.target", df.schema)
+    table = catalog.create_table(identifier, df.schema)
 
     table.append(df)
 
@@ -210,14 +219,15 @@ def test_merge_scenario_date_as_key(catalog_conn: InMemoryCatalog) -> None:
 
     assert_upsert_result(res, expected_updated, expected_inserted)
 
-    catalog.drop_table(f"{_TEST_NAMESPACE}.target")
-
 
-def test_merge_scenario_string_as_key(catalog_conn: InMemoryCatalog) -> None:
+def test_merge_scenario_string_as_key(catalog: Catalog) -> None:
     """
     tests a single insert and update; primary key is a string column
     """
 
+    identifier = "default.test_merge_scenario_string_as_key"
+    _drop_table(catalog, identifier)
+
     ctx = SessionContext()
 
     df = ctx.sql("""
@@ -226,8 +236,7 @@ def test_merge_scenario_string_as_key(catalog_conn: InMemoryCatalog) -> None:
         select 'def' as order_id, 'A' as order_type
     """).to_arrow_table()
 
-    catalog = catalog_conn
-    table = catalog.create_table(f"{_TEST_NAMESPACE}.target", df.schema)
+    table = catalog.create_table(identifier, df.schema)
 
     table.append(df)
 
@@ -246,18 +255,18 @@ def test_merge_scenario_string_as_key(catalog_conn: InMemoryCatalog) -> None:
 
     assert_upsert_result(res, expected_updated, expected_inserted)
 
-    catalog.drop_table(f"{_TEST_NAMESPACE}.target")
 
-
-def test_merge_scenario_composite_key(catalog_conn: InMemoryCatalog) -> None:
+def test_merge_scenario_composite_key(catalog: Catalog) -> None:
     """
     tests merging 200 rows with a composite key
     """
 
+    identifier = "default.test_merge_scenario_composite_key"
+    _drop_table(catalog, identifier)
+
     ctx = SessionContext()
 
-    catalog = catalog_conn
-    table = gen_target_iceberg_table(1, 200, True, ctx, catalog, _TEST_NAMESPACE)
+    table = gen_target_iceberg_table(1, 200, True, ctx, catalog, identifier)
     source_df = gen_source_dataset(101, 300, True, False, ctx)
 
     res = table.upsert(df=source_df, join_cols=["order_id", "order_line_id"])
@@ -267,43 +276,41 @@ def test_merge_scenario_composite_key(catalog_conn: InMemoryCatalog) -> None:
 
     assert_upsert_result(res, expected_updated, expected_inserted)
 
-    catalog.drop_table(f"{_TEST_NAMESPACE}.target")
 
-
-def test_merge_source_dups(catalog_conn: InMemoryCatalog) -> None:
+def test_merge_source_dups(catalog: Catalog) -> None:
     """
     tests duplicate rows in source
     """
 
+    identifier = "default.test_merge_source_dups"
+    _drop_table(catalog, identifier)
+
     ctx = SessionContext()
 
-    catalog = catalog_conn
-    table = gen_target_iceberg_table(1, 10, False, ctx, catalog, _TEST_NAMESPACE)
+    table = gen_target_iceberg_table(1, 10, False, ctx, catalog, identifier)
     source_df = gen_source_dataset(5, 15, False, True, ctx)
 
     with pytest.raises(Exception, match="Duplicate rows found in source dataset based on the key columns. No upsert executed"):
         table.upsert(df=source_df, join_cols=["order_id"])
 
-    catalog.drop_table(f"{_TEST_NAMESPACE}.target")
-
 
-def test_key_cols_misaligned(catalog_conn: InMemoryCatalog) -> None:
+def test_key_cols_misaligned(catalog: Catalog) -> None:
     """
     tests join columns missing from one of the tables
     """
 
+    identifier = "default.test_key_cols_misaligned"
+    _drop_table(catalog, identifier)
+
     ctx = SessionContext()
 
     df = ctx.sql("select 1 as order_id, date '2021-01-01' as order_date, 'A' as order_type").to_arrow_table()
 
-    catalog = catalog_conn
-    table = catalog.create_table(f"{_TEST_NAMESPACE}.target", df.schema)
+    table = catalog.create_table(identifier, df.schema)
 
     table.append(df)
 
     df_src = ctx.sql("select 1 as item_id, date '2021-05-01' as order_date, 'B' as order_type").to_arrow_table()
 
     with pytest.raises(Exception, match=r"""Field ".*" does not exist in schema"""):
         table.upsert(df=df_src, join_cols=["order_id"])
-
-    catalog.drop_table(f"{_TEST_NAMESPACE}.target")