rambleraptor commented on code in PR #3862:
URL: https://github.com/apache/iceberg-python/pull/3862#discussion_r4429467479
##########
tests/table/test_upsert.py:
##########
@@ -446,6 +446,240 @@ 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_struct_whose_fields_are_in_another_order()
-> None:
+ """The fields of a struct are compared by position, so the two have to be
in the same order.
+
+ They are: two structs of the same fields in another order are not of the
same type, and the
+ cast that follows puts the fields of the source in the order of the target.
+ """
+ source_type = pa.struct([pa.field("a", pa.int64()), pa.field("b",
pa.string())])
+ target_type = pa.struct([pa.field("b", pa.string()), pa.field("a",
pa.int64())])
+ source = pa.table({"id": pa.array([1, 2]), "nested": pa.array([{"a": 1,
"b": "x"}, {"a": 2, "b": "y"}], type=source_type)})
+ target = pa.table({"id": pa.array([1, 2]), "nested": pa.array([{"a": 1,
"b": "x"}, {"a": 9, "b": "y"}], type=target_type)})
+
+ assert get_rows_to_update(source, target, ["id"]).column("id").to_pylist()
== [2]
+
+
+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)
Review Comment:
Is it possible for us to not monkeypatch here? I'd love a more direct way to
count this, especially since this is a private API that's bound to change at
some point.
##########
pyiceberg/table/upsert_util.py:
##########
@@ -53,17 +53,91 @@ 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 field in source_column.type:
+ changed = pc.or_(
+ changed,
+ _get_changed_mask(pc.struct_field(source_column, field.name),
pc.struct_field(target_column, field.name)),
+ )
+
+ 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.
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:
Review Comment:
By moving this up here, a ValueError will be thrown instead of an empty
table being returned. I'll leave that up to the maintainers to decide if that's
alright
(My two cents: maybe try to put this below line 145?)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
