Repository: edgedb/edgedb

Issue:
Better optimization for with-bound groups
Currently we have optimizations that will sink shapes and for loop bodies into the inside of our internal group statement.
These optimizations operate on the edgeql AST, not on IR.

A big problem here is that as a result, these optimizations don't apply when the `group` is bound with `with`, which seems a fairly natural thing to want to do. Also, unfortunately, the non-optimized path is buggier than the optimized path (since it leans much more heavily on materialization). See https://github.com/edgedb/edgedb/issues/4240, https://github.com/edgedb/edgedb/issues/4241.

-----
I haven't thought about this deeply yet, but the options I have thought of so far are:
1. Do some basic analysis to see how many times a with-bound group is used, and inline it when it is only one. This would require some care since the scope might be different at the in-lining point. We can probably work around that in either a principled way or a hacky hyper-local way.
2. Change the optimizations to operate at the IR level. Probably actually this still requires doing analysis to count the number of uses



Discussion:

Resolving patch:
diff --git a/edb/edgeql/desugar_group.py b/edb/edgeql/desugar_group.py
--- a/edb/edgeql/desugar_group.py
+++ b/edb/edgeql/desugar_group.py
@@ -27,6 +27,7 @@
 
 from typing import *
 
+from edb.common import ast
 from edb.common import ordered
 from edb.common.compiler import AliasGenerator
 
@@ -165,6 +166,20 @@ def rewrite(el: qlast.GroupingElement) -> qlast.GroupingElement:
     )
 
 
+def _count_alias_uses(
+    node: qlast.Expr,
+    alias: str,
+) -> int:
+    uses = 0
+    for child in ast.find_children(node, qlast.Path):
+        match child:
+            case qlast.Path(steps=[
+                qlast.ObjectRef(module=None, name=alias2)
+            ]) if alias == alias2:
+                uses += 1
+    return uses
+
+
 def try_group_rewrite(
     node: qlast.Query,
     aliases: AliasGenerator,
@@ -199,7 +214,40 @@ def try_group_rewrite(
         )
     """
 
-    # TODO: would Python's new pattern matching fit well here???
+    # Inline trivial uses of aliases bound to a group and then
+    # immediately used, so that we can apply the other optimizations.
+    match node:
+        case qlast.SelectQuery(
+            aliases=[
+                *_,
+                qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)
+            ] as qaliases,
+            result=qlast.Shape(
+                expr=qlast.Path(steps=[
+                    qlast.ObjectRef(module=None, name=alias2)
+                ]),
+                elements=elements,
+            ) as result,
+        ) if alias == alias2 and _count_alias_uses(result, alias) == 1:
+            node = node.replace(
+                aliases=qaliases[:-1],
+                result=qlast.Shape(expr=grp, elements=elements),
+            )
+
+        case qlast.ForQuery(
+            aliases=[
+                *_,
+                qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)
+            ] as qaliases,
+            iterator=qlast.Path(steps=[
+                qlast.ObjectRef(module=None, name=alias2)
+            ]),
+            result=result,
+        ) if alias == alias2 and _count_alias_uses(result, alias) == 0:
+            node = node.replace(
+                aliases=qaliases[:-1],
+                iterator=grp,
+            )
 
     # Sink shapes into the GROUP
     if (


Test patch:
diff --git a/tests/test_edgeql_sql_codegen.py b/tests/test_edgeql_sql_codegen.py
--- a/tests/test_edgeql_sql_codegen.py
+++ b/tests/test_edgeql_sql_codegen.py
@@ -209,6 +209,20 @@ def test_codegen_group_simple(self):
             "group has unnecessary uuid_generate",
         )
 
+    def test_codegen_group_binding(self):
+        sql = self._compile('''
+        with g := (group Issue by .status)
+        select g {
+            name := .key.status.name,
+            num := count(.elements),
+        } order by .name
+        ''')
+
+        self.assertNotIn(
+            "array_agg", sql,
+            "group has unnecessary array_agg",
+        )
+
     def test_codegen_in_array_unpack_no_dupe(self):
         sql = self._compile('''
             select 1 in array_unpack(