diff --git a/pyiceberg/table/upsert_util.py b/pyiceberg/table/upsert_util.py index 6f32826eb0..900464d75f 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -53,6 +53,69 @@ 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 +# How many values of a column are turned into Python objects at a time, when PyArrow cannot +# compare them itself +_PYTHON_COMPARISON_SLICE = 10_000 + + +def _get_changed_struct_mask(source_column: pa.ChunkedArray, target_column: pa.ChunkedArray) -> pa.ChunkedArray: + """Compare two struct columns field by field, which PyArrow can do even though it cannot compare the structs.""" + # `struct_field` carries the null of the struct into its fields, so the fields of two null + # structs compare equal and only the struct itself decides for those rows + changed = pc.not_equal(pc.is_null(source_column), pc.is_null(target_column)) + + for index in range(source_column.type.num_fields): + changed = pc.or_(changed, _get_changed_mask(pc.struct_field(source_column, index), pc.struct_field(target_column, index))) + + return changed + + +def _get_changed_mask(source_column: pa.ChunkedArray, target_column: pa.ChunkedArray) -> pa.ChunkedArray: + """Return a boolean mask that flags the positions where the two columns differ, treating two nulls as equal.""" + try: + differs = pc.not_equal(source_column, target_column) + except (pa.ArrowNotImplementedError, pa.ArrowInvalid): + # PyArrow cannot compare columns with complex types + # See: https://github.com/apache/arrow/issues/35785 + if pa.types.is_struct(source_column.type) and source_column.type == target_column.type: + return _get_changed_struct_mask(source_column, target_column) + + # Two columns PyArrow refuses to compare may still hold the same values in another type: + # a naive timestamp against a zoned one, or the string of a dataframe against the + # large_string a scan reads. Comparing those in Python would call every row changed, on + # every run, and would leave a struct out of the comparison by field above. The types have + # to differ for this to make progress, the cast leaves them equal and the next round + # settles it one way or the other + if source_column.type != target_column.type: + try: + return _get_changed_mask(source_column.cast(target_column.type), target_column) + except pa.ArrowException: + # Whatever PyArrow makes of the cast, the comparison in Python below still holds + pass + + # A list or a map is left to be compared in Python, value by value. A slice at a time, + # so that the objects of a whole column are never held at once + return pa.chunked_array( + [ + [ + source_val != target_val + for source_val, target_val in zip( + source_column.slice(offset, _PYTHON_COMPARISON_SLICE).to_pylist(), + target_column.slice(offset, _PYTHON_COMPARISON_SLICE).to_pylist(), + strict=True, + ) + ] + for offset in range(0, len(source_column), _PYTHON_COMPARISON_SLICE) + ] + or [[]], + type=pa.bool_(), + ) + + # `not_equal` is null as soon as either side is null, and a null differs from a value + # but not from another null + return pc.fill_null(differs, pc.not_equal(pc.is_null(source_column), pc.is_null(target_column))) + + 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. @@ -60,10 +123,18 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols 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) + source_columns, target_columns = set(source_table.column_names), set(target_table.column_names) + if source_columns != target_columns: + raise ValueError( + f"Source table's field names are not matching the target's field names, " + f"missing: {sorted(target_columns - source_columns)}, " + f"unexpected: {sorted(source_columns - target_columns)}" + ) - non_key_cols = list(all_columns - join_cols_set) + # Kept in the order of the source rather than taken from a set difference, whose order + # varies from one process to the next + join_cols_set = set(join_cols) + non_key_cols = [col for col in source_table.column_names if col not in join_cols_set] if has_duplicate_rows(target_table, join_cols): raise ValueError("Target table has duplicate rows, aborting upsert") @@ -72,10 +143,6 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols # 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" @@ -86,39 +153,38 @@ def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols ) 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 + # Only the join columns are cast, so the width of the table does not weigh on the join # See: https://github.com/apache/arrow/issues/37542 source_index = ( - source_table.cast(target_table.schema) - .select(join_cols_set) + source_table.select(join_cols) + .cast(target_table.select(join_cols).schema) .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)))) + target_index = target_table.select(join_cols).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: + matching_indices = source_index.join(target_index, keys=join_cols, join_type="inner") + + if len(matching_indices) == 0: return source_table.schema.empty_table() + + source_indices = matching_indices[SOURCE_INDEX_COLUMN_NAME] + target_indices = matching_indices[TARGET_INDEX_COLUMN_NAME] + + # Step 4: Compare the matched rows one column at a time. Comparing them cell by cell instead + # would allocate a PyArrow scalar per cell, which does not fit in memory on a wide table. + changed = pa.chunked_array([pa.repeat(False, len(matching_indices))]) + for col in non_key_cols: + changed = pc.or_( + changed, + _get_changed_mask(source_table.column(col).take(source_indices), target_table.column(col).take(target_indices)), + ) + # Once every matched row has changed, the columns that are left cannot add anything, and + # asking is far cheaper than taking and comparing them + if pc.all(changed).as_py(): + break + + # Step 5: Take rows from source table using the indices + return source_table.take(source_indices.filter(changed)) diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..e9f10191b5 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -29,9 +29,9 @@ from pyiceberg.io.pyarrow import schema_to_pyarrow from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema -from pyiceberg.table import Table, UpsertResult +from pyiceberg.table import Table, UpsertResult, upsert_util from pyiceberg.table.snapshots import Operation -from pyiceberg.table.upsert_util import create_match_filter +from pyiceberg.table.upsert_util import create_match_filter, get_rows_to_update from pyiceberg.transforms import DayTransform from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType from tests.catalog.test_base import InMemoryCatalog @@ -446,6 +446,226 @@ def test_create_match_filter_single_condition() -> None: ) +def test_get_rows_to_update_compares_nulls() -> None: + """A null is only unchanged when the other side is null as well.""" + schema = pa.schema([pa.field("id", pa.int32()), pa.field("value", pa.string())]) + target = pa.Table.from_pylist( + [ + {"id": 1, "value": None}, + {"id": 2, "value": "b"}, + {"id": 3, "value": None}, + {"id": 4, "value": "d"}, + ], + schema=schema, + ) + source = pa.Table.from_pylist( + [ + {"id": 1, "value": None}, # null to null, unchanged + {"id": 2, "value": None}, # value to null, updated + {"id": 3, "value": "c"}, # null to value, updated + {"id": 4, "value": "d"}, # value to value, unchanged + ], + schema=schema, + ) + + assert get_rows_to_update(source, target, ["id"]).sort_by("id") == pa.Table.from_pylist( + [ + {"id": 2, "value": None}, + {"id": 3, "value": "c"}, + ], + schema=schema, + ) + + +def test_get_rows_to_update_when_a_column_is_missing() -> None: + """A source that does not carry every target column would silently null the ones it omits.""" + schema = pa.schema([pa.field("id", pa.int32()), pa.field("value", pa.string())]) + target = pa.Table.from_pylist([{"id": 1, "value": "a"}], schema=schema) + source = pa.Table.from_pylist([{"id": 1}], schema=pa.schema([pa.field("id", pa.int32())])) + + with pytest.raises(ValueError, match="field names are not matching"): + get_rows_to_update(source, target, ["id"]) + + +def test_get_rows_to_update_compares_columns_of_a_different_type() -> None: + """PyArrow refuses to compare a naive timestamp with a zoned one, the source is cast to compare. + + Comparing the two in Python would call the row changed on every run, though it holds the + same instant. + """ + target = pa.table({"id": pa.array([1, 2], pa.int32()), "ts": pa.array([1, 2], pa.timestamp("us", tz="UTC"))}) + source = pa.table({"id": pa.array([1, 2], pa.int32()), "ts": pa.array([1, 9], pa.timestamp("us"))}) + + assert get_rows_to_update(source, target, ["id"]).column("id").to_pylist() == [2] + + +def test_get_rows_to_update_when_the_cast_is_refused() -> None: + """A cast PyArrow refuses leaves the comparison in Python rather than raising.""" + source = pa.table( + { + "id": pa.array([1], pa.int32()), + "value": pa.array([[("a", 1)]], type=pa.map_(pa.string(), pa.int32())), + } + ) + target = pa.table( + { + "id": pa.array([1], pa.int32()), + "value": pa.array([[{"a": 1}]], type=pa.list_(pa.struct([("a", pa.int32())]))), + } + ) + + assert len(get_rows_to_update(source, target, ["id"])) == 1 + + +def test_get_rows_to_update_compares_a_struct_that_is_null() -> None: + """The fields of a null struct hold no meaningful value, so the struct itself decides.""" + schema = pa.schema( + [ + pa.field("id", pa.int32()), + pa.field("nested", pa.struct([pa.field("a", pa.int32()), pa.field("b", pa.string())])), + ] + ) + target = pa.Table.from_pylist( + [ + {"id": 1, "nested": None}, + {"id": 2, "nested": None}, + {"id": 3, "nested": {"a": 1, "b": "x"}}, + {"id": 4, "nested": {"a": 1, "b": "x"}}, + {"id": 5, "nested": {"a": 1, "b": None}}, + ], + schema=schema, + ) + source = pa.Table.from_pylist( + [ + {"id": 1, "nested": None}, # unchanged + {"id": 2, "nested": {"a": 1, "b": "x"}}, # null to a struct + {"id": 3, "nested": None}, # a struct to null + {"id": 4, "nested": {"a": 1, "b": "y"}}, # one field differs + {"id": 5, "nested": {"a": 1, "b": None}}, # unchanged, with a null field + ], + schema=schema, + ) + + assert [row["id"] for row in get_rows_to_update(source, target, ["id"]).sort_by("id").to_pylist()] == [2, 3, 4] + + +def test_get_rows_to_update_compares_a_nested_struct() -> None: + """A struct holding a struct is compared by recursing into the fields of both.""" + schema = pa.schema( + [ + pa.field("id", pa.int32()), + pa.field("outer", pa.struct([pa.field("inner", pa.struct([pa.field("value", pa.int32())]))])), + ] + ) + target = pa.Table.from_pylist( + [ + {"id": 1, "outer": {"inner": {"value": 1}}}, + {"id": 2, "outer": {"inner": {"value": 1}}}, + {"id": 3, "outer": {"inner": None}}, + {"id": 4, "outer": {"inner": {"value": 1}}}, + ], + schema=schema, + ) + source = pa.Table.from_pylist( + [ + {"id": 1, "outer": {"inner": {"value": 1}}}, # unchanged + {"id": 2, "outer": {"inner": {"value": 2}}}, # the innermost value differs + {"id": 3, "outer": {"inner": {"value": 1}}}, # a null inner struct to a struct + {"id": 4, "outer": {"inner": None}}, # a struct to a null inner struct + ], + schema=schema, + ) + + assert [row["id"] for row in get_rows_to_update(source, target, ["id"]).sort_by("id").to_pylist()] == [2, 3, 4] + + +def test_get_rows_to_update_compares_a_list_column() -> None: + """PyArrow cannot compare list columns either, and unlike a struct they have no fields to compare.""" + schema = pa.schema([pa.field("id", pa.int32()), pa.field("values", pa.list_(pa.int32()))]) + target = pa.Table.from_pylist( + [ + {"id": 1, "values": [1, 2]}, + {"id": 2, "values": [1, 2]}, + {"id": 3, "values": None}, + {"id": 4, "values": [1, 2]}, + ], + schema=schema, + ) + source = pa.Table.from_pylist( + [ + {"id": 1, "values": [1, 2]}, # unchanged + {"id": 2, "values": [2, 1]}, # same values, different order + {"id": 3, "values": [1]}, # null to a list + {"id": 4, "values": None}, # a list to null + ], + schema=schema, + ) + + assert [row["id"] for row in get_rows_to_update(source, target, ["id"]).sort_by("id").to_pylist()] == [2, 3, 4] + + +def test_get_rows_to_update_compares_a_column_at_a_time(monkeypatch: pytest.MonkeyPatch) -> None: + """The comparison has to walk the columns: walking the cells costs a PyArrow round trip each.""" + num_columns = 12 + comparisons = [] + + compare = upsert_util._get_changed_mask + + def record(source: pa.ChunkedArray, target: pa.ChunkedArray) -> pa.ChunkedArray: + comparisons.append((len(source), len(target))) + return compare(source, target) + + monkeypatch.setattr(upsert_util, "_get_changed_mask", record) + + for num_rows in (50, 500): + comparisons.clear() + table = pa.table({"pk": pa.array(range(num_rows)), **{f"col_{i}": pa.array([i] * num_rows) for i in range(num_columns)}}) + + assert len(get_rows_to_update(table, table, ["pk"])) == 0 + # One comparison per column, whatever the number of rows behind them + assert len(comparisons) == num_columns + + +def test_get_rows_to_update_compares_a_struct_by_its_fields(monkeypatch: pytest.MonkeyPatch) -> None: + """A struct cannot go through the comparison in Python, which builds a dict per row and column. + + The two sides carry the types the upsert path gives them: a scan reads a string as a + large_string, so the type of the target struct does not match the one of the dataframe. + """ + by_field = [] + compare_struct = upsert_util._get_changed_struct_mask + + def record(source: pa.ChunkedArray, target: pa.ChunkedArray) -> pa.ChunkedArray: + by_field.append(source.type) + return compare_struct(source, target) + + monkeypatch.setattr(upsert_util, "_get_changed_struct_mask", record) + + source = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "nested": pa.array([{"a": "x"}, {"a": "y"}], type=pa.struct([pa.field("a", pa.string())])), + } + ) + target = pa.table( + { + "id": pa.array([1, 2], pa.int32()), + "nested": pa.array([{"a": "x"}, {"a": "z"}], type=pa.struct([pa.field("a", pa.large_string())])), + } + ) + + assert get_rows_to_update(source, target, ["id"]).column("id").to_pylist() == [2] + assert by_field + + +def test_get_rows_to_update_without_any_match() -> None: + schema = pa.schema([pa.field("id", pa.int32()), pa.field("value", pa.string())]) + target = pa.Table.from_pylist([{"id": 1, "value": "a"}], schema=schema) + source = pa.Table.from_pylist([{"id": 2, "value": "b"}], schema=schema) + + assert get_rows_to_update(source, target, ["id"]) == schema.empty_table() + + def test_upsert_with_duplicate_rows_in_table(catalog: Catalog) -> None: identifier = "default.test_upsert_with_duplicate_rows_in_table"