Move build_balanced_tree to Or constructor

Author: koenvo

pyiceberg/expressions/__init__.py

@@ -21,8 +21,10 @@
 from functools import cached_property, reduce
 from typing import (
     Any,
+    Callable,
     Generic,
     Iterable,
+    Sequence,
     Set,
     Tuple,
     Type,
@@ -79,6 +81,45 @@ def __or__(self, other: BooleanExpression) -> BooleanExpression:
         return Or(self, other)
 
 
+def _build_balanced_tree(
+    operator_: Callable[[BooleanExpression, BooleanExpression], BooleanExpression], items: Sequence[BooleanExpression]
+) -> BooleanExpression:
+    """
+    Recursively constructs a balanced binary tree of BooleanExpressions using the provided binary operator.
+
+    This function is a safer and more scalable alternative to:
+        reduce(operator_, items)
+
+    Using `reduce` creates a deeply nested, unbalanced tree (e.g., operator_(a, operator_(b, operator_(c, ...)))),
+    which grows linearly with the number of items. This can lead to RecursionError exceptions in Python
+    when the number of expressions is large (e.g., >1000).
+
+    In contrast, this function builds a balanced binary tree with logarithmic depth (O(log n)),
+    helping avoid recursion issues and ensuring that expression trees remain stable, predictable,
+    and safe to traverse — especially in tools like PyIceberg that operate on large logical trees.
+
+    Parameters:
+        operator_ (Callable): A binary operator function (e.g., pyiceberg.expressions.Or, And) that takes two
+            BooleanExpressions and returns a combined BooleanExpression.
+        items (Sequence[BooleanExpression]): A sequence of BooleanExpression objects to combine.
+
+    Returns:
+        BooleanExpression: The balanced combination of all input BooleanExpressions.
+
+    Raises:
+        ValueError: If the input sequence is empty.
+    """
+    if not items:
+        raise ValueError("No expressions to combine")
+    if len(items) == 1:
+        return items[0]
+    mid = len(items) // 2
+
+    left = _build_balanced_tree(operator_, items[:mid])
+    right = _build_balanced_tree(operator_, items[mid:])
+    return operator_(left, right)
+
+
 class Term(Generic[L], ABC):
     """A simple expression that evaluates to a value."""
 
@@ -257,7 +298,7 @@ class Or(BooleanExpression):
 
     def __new__(cls, left: BooleanExpression, right: BooleanExpression, *rest: BooleanExpression) -> BooleanExpression:  # type: ignore
         if rest:
-            return reduce(Or, (left, right, *rest))
+            return _build_balanced_tree(Or, (left, right, *rest))
         if left is AlwaysTrue() or right is AlwaysTrue():
             return AlwaysTrue()
         elif left is AlwaysFalse():

pyiceberg/table/upsert_util.py

@@ -16,7 +16,6 @@
 # under the License.
 import functools
 import operator
-from typing import Callable, List, TypeVar
 
 import pyarrow as pa
 from pyarrow import Table as pyarrow_table
@@ -27,45 +26,9 @@
     BooleanExpression,
     EqualTo,
     In,
+    Or,
 )
 
-T = TypeVar("T")
-
-
-def build_balanced_tree(operator_: Callable[[T, T], T], items: List[T]) -> T:
-    """
-    Recursively constructs a balanced binary tree of expressions using the provided binary operator.
-
-    This function is a safer and more scalable alternative to:
-        reduce(operator_, items)
-
-    Using reduce creates a deeply nested, unbalanced tree (e.g., operator_(a, operator_(b, operator_(c, ...)))),
-    which grows linearly with the number of items. This can lead to RecursionError exceptions in Python
-    when the number of expressions is large (e.g., >1000).
-
-    In contrast, this function builds a balanced binary tree with logarithmic depth (O(log n)),
-    helping avoid recursion issues and ensuring that expression trees remain stable, predictable,
-    and safe to traverse — especially in tools like PyIceberg that operate on large logical trees.
-
-    Parameters:
-        operator_ (Callable[[T, T], T]): A binary operator function (e.g., pyiceberg.expressions.Or, And).
-        items (List[T]): A list of expression objects to combine.
-
-    Returns:
-        T: An expression object representing the balanced combination of all input expressions.
-
-    Raises:
-        ValueError: If the input list is empty.
-    """
-    if not items:
-        raise ValueError("No expressions to combine")
-    if len(items) == 1:
-        return items[0]
-    mid = len(items) // 2
-    left = build_balanced_tree(operator_, items[:mid])
-    right = build_balanced_tree(operator_, items[mid:])
-    return operator_(left, right)
-
 
 def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpression:
     unique_keys = df.select(join_cols).group_by(join_cols).aggregate([])
@@ -77,7 +40,12 @@ def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpre
             functools.reduce(operator.and_, [EqualTo(col, row[col]) for col in join_cols]) for row in unique_keys.to_pylist()
         ]
 
-        return AlwaysFalse() if len(filters) == 0 else build_balanced_tree(operator.or_, filters)
+        if len(filters) == 0:
+            return AlwaysFalse()
+        elif len(filters) == 1:
+            return filters[0]
+        else:
+            return Or(*filters)
 
 
 def has_duplicate_rows(df: pyarrow_table, join_cols: list[str]) -> bool:

tests/expressions/test_expressions.py

@@ -578,7 +578,7 @@ def test_negate(lhs: BooleanExpression, rhs: BooleanExpression) -> None:
         ),
         (
             Or(ExpressionA(), ExpressionB(), ExpressionA()),
-            Or(Or(ExpressionA(), ExpressionB()), ExpressionA()),
+            Or(ExpressionA(), Or(ExpressionB(), ExpressionA())),
         ),
         (Not(Not(ExpressionA())), ExpressionA()),
     ],

tests/expressions/test_visitors.py

@@ -230,9 +230,9 @@ def test_boolean_expression_visitor() -> None:
         "NOT",
         "OR",
         "EQUALTO",
-        "OR",
         "NOTEQUALTO",
         "OR",
+        "OR",
         "EQUALTO",
         "NOT",
         "AND",
@@ -408,28 +408,28 @@ def test_and_expression_binding(
                 ),
             ),
             Or(
+                BoundIn(
+                    BoundReference(
+                        field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False),
+                        accessor=Accessor(position=0, inner=None),
+                    ),
+                    {literal("bar"), literal("baz")},
+                ),
                 Or(
                     BoundIn(
                         BoundReference(
                             field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False),
                             accessor=Accessor(position=0, inner=None),
                         ),
-                        {literal("bar"), literal("baz")},
+                        {literal("bar")},
                     ),
                     BoundIn(
                         BoundReference(
                             field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False),
                             accessor=Accessor(position=0, inner=None),
                         ),
-                        {literal("bar")},
-                    ),
-                ),
-                BoundIn(
-                    BoundReference(
-                        field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False),
-                        accessor=Accessor(position=0, inner=None),
+                        {literal("baz")},
                     ),
-                    {literal("baz")},
                 ),
             ),
         ),

tests/io/test_pyarrow_visitor.py

@@ -836,5 +836,5 @@ def test_expression_to_complementary_pyarrow(
     # Notice an isNan predicate on a str column is automatically converted to always false and removed from Or and thus will not appear in the pc.expr.
     assert (
         repr(result)
-        == """<pyarrow.compute.Expression (((invert((((((string_field == "hello") and (float_field > 100)) or (is_nan(float_field) and (double_field == 0))) or (float_field > 100)) and invert(is_null(double_field, {nan_is_null=false})))) or is_null(float_field, {nan_is_null=false})) or is_null(string_field, {nan_is_null=false})) or is_nan(double_field))>"""
+        == """<pyarrow.compute.Expression (((invert(((((string_field == "hello") and (float_field > 100)) or ((is_nan(float_field) and (double_field == 0)) or (float_field > 100))) and invert(is_null(double_field, {nan_is_null=false})))) or is_null(float_field, {nan_is_null=false})) or is_null(string_field, {nan_is_null=false})) or is_nan(double_field))>"""
     )