Fix PyArrow stack-overflow segfault in PyIceberg's multi-column upsert match filter

pyiceberg/table/upsert_util.py

@@ -23,6 +23,7 @@
 
 from pyiceberg.expressions import (
     AlwaysFalse,
+    And,
     BooleanExpression,
     EqualTo,
     In,
@@ -33,19 +34,40 @@
 def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpression:
     unique_keys = df.select(join_cols).group_by(join_cols).aggregate([])
 
+    if unique_keys.num_rows == 0:
+        return AlwaysFalse()
+
     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)
+        return In(join_cols[0], unique_keys.column(join_cols[0]).to_pylist())
+
+    # Fold the column that leaves the fewest distinct "prefix" combinations into
+    # an In(); this minimises the disjunct count regardless of column order.
+    in_col = min(
+        join_cols,
+        key=lambda cand: unique_keys.select([c for c in join_cols if c != cand])
+        .group_by([c for c in join_cols if c != cand])
+        .aggregate([])
+        .num_rows,
+    )
+    prefix_cols = [c for c in join_cols if c != in_col]
+
+    # The group keys come first (in prefix_cols order) followed by the list aggregate.
+    # Rename the aggregate to a sentinel so it cannot collide with a join column that
+    # happens to be named f"{in_col}_list".
+    in_values_col = "__in_values"
+    while in_values_col in prefix_cols:
+        in_values_col += "_"
+    grouped = unique_keys.group_by(prefix_cols).aggregate([(in_col, "list")]).rename_columns([*prefix_cols, in_values_col])
+
+    disjuncts: list[BooleanExpression] = []
+    for row in grouped.to_pylist():
+        eqs = [EqualTo(c, row[c]) for c in prefix_cols]
+        prefix_pred = functools.reduce(operator.and_, eqs) if len(eqs) > 1 else eqs[0]
+        disjuncts.append(And(prefix_pred, In(in_col, row[in_values_col])))
+
+    if len(disjuncts) == 1:
+        return disjuncts[0]
+    return Or(*disjuncts)
 
 
 def has_duplicate_rows(df: pyarrow_table, join_cols: list[str]) -> bool:

tests/table/test_upsert.py

@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import itertools
 from datetime import datetime
 from pathlib import PosixPath
 
@@ -24,15 +25,17 @@
 
 from pyiceberg.catalog import Catalog
 from pyiceberg.exceptions import NoSuchTableError
-from pyiceberg.expressions import AlwaysTrue, And, EqualTo, Reference
+from pyiceberg.expressions import AlwaysFalse, AlwaysTrue, And, BooleanExpression, EqualTo, In, Or, Reference
 from pyiceberg.expressions.literals import LongLiteral
+from pyiceberg.expressions.visitors import expression_evaluator
 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.snapshots import Operation
 from pyiceberg.table.upsert_util import create_match_filter
 from pyiceberg.transforms import DayTransform
+from pyiceberg.typedef import Record
 from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType
 from tests.catalog.test_base import InMemoryCatalog
 
@@ -440,10 +443,228 @@ 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"])
-    assert expr == And(
-        EqualTo(term=Reference(name="order_id"), literal=LongLiteral(101)),
-        EqualTo(term=Reference(name="order_line_id"), literal=LongLiteral(1)),
+    # Be insensitive to left/right operands
+    op1 = EqualTo(term=Reference(name="order_id"), literal=LongLiteral(101))
+    op2 = EqualTo(term=Reference(name="order_line_id"), literal=LongLiteral(1))
+    assert expr == And(op1, op2) or expr == And(op2, op1)
+
+
+def _assert_match_filter_selects(data: list[dict[str, int]], join_cols: list[str], schema: Schema) -> BooleanExpression:
+    """Assert the filter from ``create_match_filter`` matches exactly the unique source keys.
+
+    Rather than asserting a specific expression tree (which is implementation-specific),
+    this binds the filter and evaluates it against the full cross-product of the values
+    observed per column. The filter must accept exactly the unique keys present in
+    ``data`` and reject every other combination, so any over- or under-matching
+    (e.g. a cross-product regression) is caught. This holds for any correct
+    implementation of ``create_match_filter``.
+    """
+    arrow_schema = schema_to_pyarrow(schema)
+    table = pa.Table.from_pylist(data, schema=arrow_schema)
+    expr = create_match_filter(table, join_cols)
+
+    field_names = [field.name for field in schema.fields]
+    expected_keys = {tuple(row[name] for name in field_names) for row in data}
+    domains = [sorted({row[name] for row in data}) for name in field_names]
+
+    evaluate = expression_evaluator(schema, expr, case_sensitive=True)
+    for candidate in itertools.product(*domains):
+        key = dict(zip(field_names, candidate, strict=True))
+        should_match = candidate in expected_keys
+        verb = "rejected matching" if should_match else "matched non-matching"
+        assert evaluate(Record(*candidate)) is should_match, f"Filter {expr} {verb} key {key}"
+
+    return expr
+
+
+def _count_disjuncts(expr: BooleanExpression) -> int:
+    """Count the leaves of a top-level disjunction."""
+    if isinstance(expr, Or):
+        return _count_disjuncts(expr.left) + _count_disjuncts(expr.right)
+    return 1
+
+
+def test_create_match_filter_single_prefix_group() -> None:
+    """
+    Test create_match_filter with multiple key columns whose rows all share a single prefix combination.
+
+    The filter must match the (one order_id, many order_line_id) keys and nothing else.
+    """
+    schema = Schema(
+        NestedField(1, "order_id", IntegerType(), required=True),
+        NestedField(2, "order_line_id", IntegerType(), required=True),
+    )
+    data = [
+        {"order_id": 101, "order_line_id": 1},
+        {"order_id": 101, "order_line_id": 2},
+        {"order_id": 101, "order_line_id": 3},
+        {"order_id": 101, "order_line_id": 3},  # duplicate
+    ]
+    _assert_match_filter_selects(data, ["order_id", "order_line_id"], schema)
+
+
+def test_create_match_filter_multiple_prefix_groups() -> None:
+    """
+    Test create_match_filter with multiple key columns that yield several distinct prefix combinations.
+
+    The filter must match exactly the listed composite keys and must NOT match cross-product
+    combinations that never appear together (e.g. order_id 101 with order_line_id 2).
+    """
+    schema = Schema(
+        NestedField(1, "order_id", IntegerType(), required=True),
+        NestedField(2, "order_line_id", IntegerType(), required=True),
+    )
+    data = [
+        {"order_id": 101, "order_line_id": 1},
+        {"order_id": 102, "order_line_id": 1},
+        {"order_id": 103, "order_line_id": 1},
+        {"order_id": 201, "order_line_id": 2},
+        {"order_id": 202, "order_line_id": 2},
+    ]
+    _assert_match_filter_selects(data, ["order_id", "order_line_id"], schema)
+
+
+@pytest.mark.parametrize("join_cols", [["order_id", "region"], ["region", "order_id"]])
+def test_create_match_filter_collapses_low_cardinality_prefix(join_cols: list[str]) -> None:
+    """Expression size depends on prefix cardinality, regardless of join-column order."""
+    schema = Schema(
+        NestedField(1, "order_id", IntegerType(), required=True),
+        NestedField(2, "region", IntegerType(), required=True),
+    )
+    data = [{"order_id": order_id, "region": order_id % 4} for order_id in range(100)]
+
+    expr = _assert_match_filter_selects(data, join_cols, schema)
+
+    assert _count_disjuncts(expr) == 4
+
+
+def test_create_match_filter_single_column() -> None:
+    """A single join column collapses to a single In() over the unique values."""
+    schema = pa.schema([pa.field("order_id", pa.int32())])
+    table = pa.Table.from_pylist([{"order_id": 1}, {"order_id": 2}, {"order_id": 2}], schema=schema)
+    assert create_match_filter(table, ["order_id"]) == In("order_id", [1, 2])
+
+
+def test_create_match_filter_single_column_single_value() -> None:
+    """A single unique value collapses the In() down to an EqualTo()."""
+    schema = pa.schema([pa.field("order_id", pa.int32())])
+    table = pa.Table.from_pylist([{"order_id": 1}, {"order_id": 1}], schema=schema)
+    assert create_match_filter(table, ["order_id"]) == EqualTo("order_id", 1)
+
+
+def test_create_match_filter_empty_input() -> None:
+    """An empty source matches nothing (AlwaysFalse), for both single and composite keys."""
+    schema = pa.schema([pa.field("order_id", pa.int32()), pa.field("order_line_id", pa.int32())])
+    empty = pa.Table.from_pylist([], schema=schema)
+    assert create_match_filter(empty, ["order_id"]) == AlwaysFalse()
+    assert create_match_filter(empty, ["order_id", "order_line_id"]) == AlwaysFalse()
+
+
+def test_create_match_filter_three_columns() -> None:
+    """
+    Test create_match_filter with three key columns.
+
+    Exercises the multi-column prefix branch where the prefix predicate is an And of two
+    EqualTo() conjuncts combined with an In() over the folded column.
+    """
+    schema = Schema(
+        NestedField(1, "a", IntegerType(), required=True),
+        NestedField(2, "b", IntegerType(), required=True),
+        NestedField(3, "c", IntegerType(), required=True),
+    )
+    data = [
+        {"a": 1, "b": 1, "c": 1},
+        {"a": 1, "b": 1, "c": 2},
+        {"a": 1, "b": 1, "c": 3},
+        {"a": 2, "b": 9, "c": 5},
+        {"a": 2, "b": 9, "c": 6},
+    ]
+    _assert_match_filter_selects(data, ["a", "b", "c"], schema)
+
+
+def test_create_match_filter_column_named_like_aggregate() -> None:
+    """
+    Regression test for #3509 review feedback.
+
+    A join column named ``<in_col>_list`` must not collide with the internal list-aggregation
+    column used to fold values into an In(). Before the fix this raised a TypeError.
+    """
+    schema = Schema(
+        NestedField(1, "a", IntegerType(), required=True),
+        NestedField(2, "a_list", IntegerType(), required=True),
+    )
+    data = [
+        {"a": 1, "a_list": 7},
+        {"a": 2, "a_list": 7},
+        {"a": 3, "a_list": 8},
+    ]
+    _assert_match_filter_selects(data, ["a", "a_list"], schema)
+
+
+def test_create_match_filter_internal_aggregate_name_collision() -> None:
+    """Internal aggregate names advance until they no longer collide with prefix columns."""
+    schema = Schema(
+        NestedField(1, "id", IntegerType(), required=True),
+        NestedField(2, "__in_values", IntegerType(), required=True),
+        NestedField(3, "__in_values_", IntegerType(), required=True),
     )
+    data = [
+        {"id": 1, "__in_values": 7, "__in_values_": 70},
+        {"id": 2, "__in_values": 7, "__in_values_": 70},
+        {"id": 3, "__in_values": 8, "__in_values_": 80},
+        {"id": 4, "__in_values": 8, "__in_values_": 80},
+    ]
+
+    expr = _assert_match_filter_selects(data, ["id", "__in_values", "__in_values_"], schema)
+
+    assert _count_disjuncts(expr) == 2
+
+
+def test_upsert_grouped_composite_key_final_state(catalog: Catalog) -> None:
+    """Grouped composite-key upserts update, preserve, and insert the exact expected rows."""
+    identifier = "default.test_upsert_grouped_composite_key_final_state"
+    _drop_table(catalog, identifier)
+
+    target_rows = 1_000
+    update_rows = 500
+    insert_rows = 100
+    schema = pa.schema(
+        [
+            pa.field("order_id", pa.int64(), nullable=False),
+            pa.field("region", pa.string(), nullable=False),
+            pa.field("amount", pa.int64(), nullable=False),
+        ]
+    )
+
+    def make(order_ids: range, amount: int) -> pa.Table:
+        return pa.Table.from_pylist(
+            [{"order_id": oid, "region": "ABCD"[oid % 4], "amount": amount} for oid in order_ids],
+            schema=schema,
+        )
+
+    tbl = catalog.create_table(identifier, schema)
+    tbl.append(make(range(target_rows), amount=1))
+
+    source = pa.concat_tables(
+        [
+            make(range(update_rows), amount=2),
+            make(range(target_rows, target_rows + insert_rows), amount=2),
+        ]
+    )
+
+    res = tbl.upsert(source, join_cols=["order_id", "region"])
+    assert res.rows_updated == update_rows
+    assert res.rows_inserted == insert_rows
+
+    actual = {row["order_id"]: {"region": row["region"], "amount": row["amount"]} for row in tbl.scan().to_arrow().to_pylist()}
+    expected = {
+        order_id: {
+            "region": "ABCD"[order_id % 4],
+            "amount": 2 if order_id < update_rows or order_id >= target_rows else 1,
+        }
+        for order_id in range(target_rows + insert_rows)
+    }
+    assert actual == expected
 
 
 def test_upsert_with_duplicate_rows_in_table(catalog: Catalog) -> None: