perf: build partition filter with balanced tree to avoid RecursionError

Author: Yingjian Wu

pyiceberg/table/__init__.py

@@ -366,21 +366,19 @@ def _build_partition_predicate(
         Returns:
             A predicate matching any of the input partition records.
         """
-        partition_fields = [schema.find_field(field.source_id).name for field in spec.fields]
-
-        expr: BooleanExpression = AlwaysFalse()
-        for partition_record in partition_records:
-            match_partition_expression: BooleanExpression = AlwaysTrue()
-
-            for pos, partition_field in enumerate(partition_fields):
-                predicate = (
-                    EqualTo(Reference(partition_field), partition_record[pos])
-                    if partition_record[pos] is not None
-                    else IsNull(Reference(partition_field))
-                )
-                match_partition_expression = And(match_partition_expression, predicate)
-            expr = Or(expr, match_partition_expression)
-        return expr
+        partition_fields = [schema.find_field(f.source_id).name for f in spec.fields]
+        if not partition_records or not partition_fields:
+            return AlwaysFalse()
+
+        def _match(record: Record) -> BooleanExpression:
+            parts: list[BooleanExpression] = [
+                EqualTo(Reference(name), record[pos]) if record[pos] is not None else IsNull(Reference(name))
+                for pos, name in enumerate(partition_fields)
+            ]
+            return And(*parts) if len(parts) > 1 else parts[0]
+
+        per_record = [_match(r) for r in partition_records]
+        return Or(*per_record) if len(per_record) > 1 else per_record[0]
 
     def _append_snapshot_producer(
         self, snapshot_properties: dict[str, str], branch: str | None = MAIN_BRANCH

tests/table/test_init.py

@@ -1753,3 +1753,19 @@ def test_check_uuid_passes_when_match(table_v2: Table, example_table_metadata_v2
     new_metadata = TableMetadataV2(**example_table_metadata_v2)
     # Should not raise with same uuid
     Table._check_uuid(table_v2.metadata, new_metadata)
+
+
+def test_build_large_partition_predicate(table_v2: Table) -> None:
+    """A left-folded Or chain over 5000 records would be depth-5000 and crash bind()
+    (Python's default recursion limit is ~1000). The balanced tree has depth ~14."""
+    from pyiceberg.expressions.visitors import bind
+    from pyiceberg.typedef import Record
+
+    with table_v2.transaction() as tx:
+        expr = tx._build_partition_predicate(
+            partition_records={Record(i) for i in range(5000)},
+            spec=table_v2.metadata.spec(),
+            schema=table_v2.metadata.schema(),
+        )
+
+    bind(table_v2.metadata.schema(), expr, case_sensitive=True)