Repository: geldata/gel
Pull request: Rebalance deeply nested `UNION` operators
We could do other associative operators if we wanted, like ++. If we
wanted to do +, we'd have to do it after type inference so that we
don't do it on floating point values, where it doesn't hold.

Fixes #3669.

Problem statement:
Bulk Inserting Error: maximum recursion depth exceeded
I'm getting the error `edgedb.errors.InternalServerError: maximum recursion depth exceeded while calling a Python object` when I'm bulk-inserting ~50 small objects in one query.

The schema is:
```
  type Person {
    required property first_name -> str;
    property last_name -> str;
    multi property tags -> str;
    property ordered_tags -> array<str>;
    required property inserted_at -> datetime {
        readonly := true;
        default := datetime_current();
    }
    property created_at -> datetime;
    multi link people_i_follow -> Person {
        on target delete allow;
    }
    required property slug -> str {
        constraint exclusive;
    }
    multi link my_followers := .<people_i_follow[is Person];
    link best_friend -> Person {
        on target delete allow;
    }
    property full_name :=
      .first_name ++ ' ' ++ .last_name
      if exists .last_name
      else .first_name;
  }
```

And the code creates this query:
```
with
model_1 := (insert Person {first_name := <str>$first_name_1, slug := <str>$slug_1}),
model_2 := (insert Person {first_name := <str>$first_name_2, slug := <str>$slug_2}),
...
model_50 := (insert Person {first_name := <str>$first_name_50, slug := <str>$slug_50}),
models := (model_1 union model_1 ... union model_50)
select models {id}
```

My current workarounds are 1) only bulk inserting a small number of objects at a time ~3, and 2) asynchronously inserting objects individually.

Hints from issue discussion:
It seems like the size of the query overflows the call stack in the compiler.

This can easily be solved by using a `for` instead:

```
for data in array_unpack(<array<tuple<first_name: str, slug: str>>>$json_data)
union (insert Person { first_name := data.first_name, slug := data.slug })
```

where `data` is a json array in this format:

```
[{"first_name": "foo", "slug": "bar"}, ...]
```
@elprans got it. What if I want to do a combination of inserts and updates for various objects? Or what if the objects I'm inserting have different fields (like some have a last name I want to insert too)?

For context, I'm making a batch function that allows any group of mutations to atomically add to the DB in bulk.

The functionality I'm going for is this:
https://firebase.google.com/docs/firestore/manage-data/transactions#batched-writes
Use individual statements for each kind of mutation or combination of fields and wrap the whole thing in a [transaction](https://www.edgedb.com/docs/clients/01_js/driver#transactions).
Got it. Is there a difference between doing a transaction and doing a for loop?
An EdgeQL `for` or a programming language `for` loop?
An EdgeQL `for`
An EdgeQL `for` is a single statement and all statements are atomic.  Transaction blocks are also atomic, but each statement within a block is sent individually in sequence.  So, from the latency standpoint, it's better to use EdgeQL `for` whenever you can.
Using `for` to do bulk operations in general is a better approach than generating queries dynamically (dynamic queries will miss the cache a lot!).

That said, as a workaround here you can use `{model_1, model_2, ..., model_50}` instead of `union`. The issue here is that `union` is a left associative binary operation, so it is really `(........((((model_1 union model_2) union model_3) union model_4.....)`

We could probably work around this by rebalancing `union` the same way we do explicit sets?
Can also repro with `select (0+0 union 1+0 union 2+0 union 3+0 union 4+0 union 5+0 union 6+0 union 7+0 union 8+0 union 9+0 union 10+0 union 11+0 union 12+0 union 13+0 union 14+0 union 15+0 union 16+0 union 17+0 union 18+0 union 19+0 union 20+0 union 21+0 union 22+0 union 23+0 union 24+0 union 25+0 union 26+0 union 27+0 union 28+0 union 29+0 union 30+0 union 31+0 union 32+0 union 33+0 union 34+0 union 35+0 union 36+0 union 37+0 union 38+0 union 39+0 union 40+0 union 41+0 union 42+0 union 43+0 union 44+0 union 45+0 union 46+0 union 47+0 union 48+0 union 49+0);`
Yes, let's rebalance the `union` tree. After all, `{}` is a syntax sugar of `union`.

Commit message: Rebalance deeply nested `UNION` operators

We could do other associative operators if we wanted, like ++. If we
wanted to do +, we'd have to do it after type inference so that we
don't do it on floating point values, where it doesn't hold.

Fixes #3669.

Patch to review:
diff --git a/edb/edgeql/ast.py b/edb/edgeql/ast.py
index e27f2deabe3..74c2160859b 100644
--- a/edb/edgeql/ast.py
+++ b/edb/edgeql/ast.py
@@ -287,6 +287,7 @@ class BinOp(Expr):
     left: Expr
     op: str
     right: Expr
+    rebalanced: bool = False
 
 
 class SetConstructorOp(BinOp):
diff --git a/edb/edgeql/compiler/expr.py b/edb/edgeql/compiler/expr.py
index 91d81a09a91..86fcebfb299 100644
--- a/edb/edgeql/compiler/expr.py
+++ b/edb/edgeql/compiler/expr.py
@@ -71,10 +71,57 @@ def compile_Path(
     return stmt.maybe_add_view(setgen.compile_path(expr, ctx=ctx), ctx=ctx)
 
 
+def _balance(
+    elements: Sequence[qlast.Expr],
+    ctor: Callable[
+        [qlast.Expr, qlast.Expr, Optional[ctx_utils.ParserContext]],
+        qlast.Expr
+    ],
+    context: Optional[ctx_utils.ParserContext],
+) -> qlast.Expr:
+    mid = len(elements) // 2
+    ls, rs = elements[:mid], elements[mid:]
+    ls_context = rs_context = None
+    if len(ls) > 1 and ls[0].context and ls[-1].context:
+        ls_context = ctx_utils.merge_context([
+            ls[0].context, ls[-1].context])
+    if len(rs) > 1 and rs[0].context and rs[-1].context:
+        rs_context = ctx_utils.merge_context([
+            rs[0].context, rs[-1].context])
+
+    return ctor(
+        (
+            _balance(ls, ctor, ls_context)
+            if len(ls) > 1 else ls[0]
+        ),
+        (
+            _balance(rs, ctor, rs_context)
+            if len(rs) > 1 else rs[0]
+        ),
+        context,
+    )
+
+
+REBALANCED_OPS = {'UNION'}
+REBALANCE_THRESHOLD = 10
+
+
 @dispatch.compile.register(qlast.SetConstructorOp)
 @dispatch.compile.register(qlast.BinOp)
 def compile_BinOp(
         expr: qlast.BinOp, *, ctx: context.ContextLevel) -> irast.Set:
+    # Rebalance some associative operations to avoid deeply nested ASTs
+    if expr.op in REBALANCED_OPS and not expr.rebalanced:
+        elements = collect_binop(expr, expr.op)
+        # Don't bother rebalancing small groups
+        if len(elements) >= REBALANCE_THRESHOLD:
+            balanced = _balance(
+                elements,
+                lambda l, r, c: qlast.BinOp(
+                    left=l, right=r, op=expr.op, rebalanced=True, context=c),
+                expr.context
+            )
+            return dispatch.compile(balanced, ctx=ctx)
 
     op_node = func.compile_operator(
         expr, op_name=expr.op, qlargs=[expr.left, expr.right], ctx=ctx)
@@ -150,32 +197,11 @@ def compile_Set(
             # Turn it into a tree of UNIONs so we only blow up the nesting
             # depth logarithmically.
             # TODO: Introduce an N-ary operation that handles the whole thing?
-            mid = len(elements) // 2
-            ls, rs = elements[:mid], elements[mid:]
-            ls_context = rs_context = None
-            if len(ls) > 1 and ls[0].context and ls[-1].context:
-                ls_context = ctx_utils.merge_context([
-                    ls[0].context, ls[-1].context])
-            if len(rs) > 1 and rs[0].context and rs[-1].context:
-                rs_context = ctx_utils.merge_context([
-                    rs[0].context, rs[-1].context])
-
-            bigunion = qlast.SetConstructorOp(
-                left=(
-                    qlast.Set(
-                        elements=ls,
-                        context=ls_context,
-                    )
-                    if len(ls) > 1 else ls[0]
-                ),
-                right=(
-                    qlast.Set(
-                        elements=rs,
-                        context=rs_context,
-                    )
-                    if len(rs) > 1 else rs[0]
-                ),
-                context=expr.context,
+            bigunion = _balance(
+                elements,
+                lambda l, r, c: qlast.SetConstructorOp(
+                    left=l, right=r, context=c),
+                expr.context
             )
             return dispatch.compile(bigunion, ctx=ctx)
     else:
@@ -666,3 +692,17 @@ def flatten_set(expr: qlast.Set) -> List[qlast.Expr]:
             elements.append(el)
 
     return elements
+
+
+def collect_binop(expr: qlast.Expr, op: str) -> List[qlast.Expr]:
+    elements = []
+    def _go(expr: qlast.Expr) -> None:
+        if isinstance(expr, qlast.BinOp) and expr.op == op:
+            _go(expr.left)
+            _go(expr.right)
+        else:
+            elements.append(expr)
+
+    _go(expr)
+
+    return elements
diff --git a/tests/test_edgeql_select.py b/tests/test_edgeql_select.py
index e06e4e79238..e67492e9c6d 100644
--- a/tests/test_edgeql_select.py
+++ b/tests/test_edgeql_select.py
@@ -6288,6 +6288,34 @@ async def test_edgeql_select_big_set_literal(self):
 
         assert len(res) == 100
 
+    async def test_edgeql_select_big_unions(self):
+        res = await self.con.query("""
+            SELECT (
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,) union
+                 (1,) union (1,) union (1,) union (1,) union (1,)
+            );
+        """)
+
+        assert len(res) == 100
+
     async def test_edgeql_select_shape_on_scalar(self):
         with self.assertRaisesRegex(
             edgedb.QueryError,


Code review comments:
diff_hunk: @@ -666,3 +692,17 @@ def flatten_set(expr: qlast.Set) -> List[qlast.Expr]:
             elements.append(el)
 
     return elements
+
+
+def collect_binop(expr: qlast.Expr, op: str) -> List[qlast.Expr]:
+    elements = []
+    def _go(expr: qlast.Expr) -> None:
path: edb/edgeql/compiler/expr.py
text: @user1:
Perhaps make this into an iterative traversal to further avoid recursion issues?

@author:
I think I had it in my head that there was some pass before this point that would do a full tree traversal, and so the important bit was just dodging the ~10x stack frame blowup in pgsql. But that's wrong, and this was the first thing hit when I made it much bigger