Fix volatile with bindings in group statements.

edb/pgsql/compiler/clauses.py

@@ -23,6 +23,7 @@
 
 import random
 
+from edb.common import topological
 from edb.common import ast as ast_visitor
 
 from edb.edgeql import qltypes
@@ -114,12 +115,34 @@ def compile_materialized_exprs(
         matctx.materializing |= {stmt}
         matctx.expr_exposed = True
 
-        # HACK: Sort longer paths before shorter ones
-        # We want foo->bar to appear before foo
-        mat_sets = sorted(
-            (stmt.materialized_sets.values()),
-            key=lambda m: -len(m.materialized.path_id),
-        )
+        # Determine the order to materialize sets. If a set A references another
+        # set B, B should materialize first.
+        mat_set_dependency_graph: dict[
+            irast.PathId,
+            topological.DepGraphEntry[
+                irast.PathId, irast.MaterializedSet, None
+            ],
+        ] = {}
+        for mat_set in stmt.materialized_sets.values():
+            path_id = mat_set.materialized.path_id
+
+            mat_set_dependency_graph[path_id] = topological.DepGraphEntry(
+                item=mat_set,
+                deps={
+                    child.path_id
+                    for child in ast_visitor.find_children(
+                        mat_set.materialized, irast.Set
+                    )
+                    if isinstance(child.expr, irast.MaterializedExpr)
+                },
+            )
+
+        mat_sets: list[irast.MaterializedSet] = [
+            mat_set
+            for mat_set in topological.sort(
+                mat_set_dependency_graph, allow_unresolved=True,
+            )
+        ]
 
         for mat_set in mat_sets:
             if len(mat_set.uses) <= 1:

edb/pgsql/compiler/group.py

@@ -177,6 +177,9 @@ def _compile_group(
 
     query = ctx.stmt
 
+    # Process materialized sets
+    clauses.compile_materialized_exprs(query, stmt, ctx=ctx)
+
     # Compile a GROUP BY into a subquery, along with all the aggregations
     with ctx.subrel() as groupctx:
         grouprel = groupctx.rel

edb/pgsql/compiler/relgen.py

@@ -1536,9 +1536,10 @@ def process_set_as_subquery(
             return _new_subquery_stmt_set_rvar(ir_set, stmt, ctx=newctx)
 
         # materialized refs should always get picked up by now
-        assert not isinstance(expr, irast.MaterializedExpr), (
-            f"Can't find materialized set {ir_set.path_id}"
-        )
+        if isinstance(expr, irast.MaterializedExpr):
+            raise AssertionError(
+                f"Can't find materialized set {ir_set.path_id}"
+            )
         assert isinstance(expr, irast.Stmt)
 
         inner_set = expr.result

tests/test_edgeql_group.py

@@ -1089,6 +1089,91 @@ async def test_edgeql_group_binding_01(self):
             ])
         )
 
+    async def test_edgeql_group_binding_02(self):
+        await self.assert_query_result(
+            '''
+                group (
+                    with X := {random(), random()}
+                    for x in enumerate(X)
+                        select {a := 1, b := x.0, rand := x.1}
+                ) by .a;
+            ''',
+            tb.bag([
+                {
+                    'key': {'a': 1},
+                    'grouping': ['a'],
+                    'elements': [
+                        {'a': 1, 'b': 0},
+                        {'a': 1, 'b': 1},
+                    ]
+                }
+            ])
+        )
+
+    async def test_edgeql_group_binding_03(self):
+        await self.assert_query_result(
+            '''
+                with X := {random(), random()}
+                group (
+                    for x in enumerate(X)
+                        select {a := 1, b := x.0, rand := x.1}
+                ) by .a;
+            ''',
+            tb.bag([
+                {
+                    'key': {'a': 1},
+                    'grouping': ['a'],
+                    'elements': [
+                        {'a': 1, 'b': 0},
+                        {'a': 1, 'b': 1},
+                    ]
+                }
+            ])
+        )
+
+    async def test_edgeql_group_binding_04(self):
+        await self.assert_query_result(
+            '''
+                with X := {random(), random()},
+                Y := (
+                    for x in enumerate(X)
+                        select {a := 1, b := x.0, rand := x.1}
+                )
+                group Y { a, b } by .a;
+            ''',
+            tb.bag([
+                {
+                    'key': {'a': 1},
+                    'grouping': ['a'],
+                    'elements': [
+                        {'a': 1, 'b': 0},
+                        {'a': 1, 'b': 1},
+                    ]
+                }
+            ])
+        )
+
+    async def test_edgeql_group_binding_05(self):
+        await self.assert_query_result(
+            '''
+                with Y := (
+                    for x in enumerate({random(), random()})
+                        select {a := 1, b := x.0, rand := x.1}
+                )
+                group Y { a, b } by .a;
+            ''',
+            tb.bag([
+                {
+                    'key': {'a': 1},
+                    'grouping': ['a'],
+                    'elements': [
+                        {'a': 1, 'b': 0},
+                        {'a': 1, 'b': 1},
+                    ]
+                }
+            ])
+        )
+
     async def test_edgeql_group_ordering_01(self):
         res = [
             {