Skip to content

Commit

Permalink
blarg
Browse files Browse the repository at this point in the history
  • Loading branch information
@dnwpark
dnwpark committed Nov 30, 2024
1 parent fb90d5f commit c83c843
Show file tree
Hide file tree
Showing 3 changed files with 218 additions and 16 deletions.
34 changes: 19 additions & 15 deletions edb/edgeql/compiler/stmt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -309,32 +309,32 @@ def compile_InternalGroupQuery(
expr: qlast.InternalGroupQuery, *, ctx: context.ContextLevel
) -> irast.Set:
# We disallow use of FOR GROUP except for when running in test mode.
if not expr.from_desugaring and not ctx.env.options.testmode:
raise errors.UnsupportedFeatureError(
"'FOR GROUP' is an internal testing feature",
span=expr.span,
)

_protect_expr(expr.subject, ctx=ctx)
_protect_expr(expr.result, ctx=ctx)

with ctx.subquery() as sctx:
stmt = irast.GroupStmt(by=expr.by)

init_stmt(stmt, expr, ctx=sctx, parent_ctx=ctx)

with sctx.newscope(fenced=True) as topctx:
init_stmt(stmt, expr, ctx=topctx, parent_ctx=ctx)

# N.B: Subject is exposed because we want any shape on the
# subject to be exposed on bare references to the group
# alias. This is frankly pretty dodgy behavior for
# FOR GROUP to have but the real GROUP needs to
# maintain shapes, and this is the easiest way to handle
# that.
stmt.subject = compile_result_clause(
expr.subject,
result_alias=expr.subject_alias,
exprtype=s_types.ExprType.Group,
ctx=topctx)
stmt.subject = setgen.scoped_set(
compile_result_clause(
(expr.subject),
result_alias=expr.subject_alias,
exprtype=s_types.ExprType.Group,
ctx=topctx,
),
ctx=topctx,
)

if topctx.partial_path_prefix:
pathctx.register_set_in_scope(
Expand Down Expand Up @@ -425,10 +425,14 @@ def compile_InternalGroupQuery(
stmt.grouping_binding, path_scope=bctx.path_scope, ctx=bctx
)

stmt.result = compile_result_clause(
astutils.ensure_ql_query(expr.result),
result_alias=expr.result_alias,
ctx=bctx)
stmt.result = setgen.scoped_set(
compile_result_clause(
astutils.ensure_ql_query(expr.result),
result_alias=expr.result_alias,
ctx=bctx,
),
ctx=bctx,
)

stmt.where = clauses.compile_where_clause(expr.where, ctx=bctx)

Expand Down
4 changes: 3 additions & 1 deletion edb/pgsql/compiler/group.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,9 @@ def _compile_group(

query = ctx.stmt

# Process materialized sets
clauses.compile_materialized_exprs(ctx.rel, stmt, ctx=ctx)

# Compile a GROUP BY into a subquery, along with all the aggregations
with ctx.subrel() as groupctx:
grouprel = groupctx.rel
Expand All @@ -197,7 +200,6 @@ def _compile_group(
ctx=groupctx,
)


# First compile the actual subject
# subrel *solely* for path id map reasons
with groupctx.subrel() as subjctx:
Expand Down
196 changes: 196 additions & 0 deletions tests/test_edgeql_internal_group.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1540,3 +1540,199 @@ async def test_edgeql_igroup_reshape_02(self):
{"avg_cost": 2, "element": "Water"},
])
)

async def test_edgeql_with_alias_simple_01(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
C := Card
FOR GROUP C
USING e := .element
BY e
IN g
UNION
{
key := e,
names := g.name,
};
""",
tb.bag([
{'key': 'Water', 'names': ['Bog monster', 'Giant turtle']},
{'key': 'Fire', 'names': ['Imp', 'Dragon']},
{'key': 'Earth', 'names': ['Dwarf', 'Golem']},
{'key': 'Air', 'names': ['Sprite', 'Giant eagle', 'Djinn']},
])
)

async def test_edgeql_with_alias_simple_02(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
C := Card
FOR GROUP (select C)
USING e := .element
BY e
IN g
UNION
{
key := e,
names := g.name,
};
""",
tb.bag([
{'key': 'Water', 'names': ['Bog monster', 'Giant turtle']},
{'key': 'Fire', 'names': ['Imp', 'Dragon']},
{'key': 'Earth', 'names': ['Dwarf', 'Golem']},
{'key': 'Air', 'names': ['Sprite', 'Giant eagle', 'Djinn']},
])
)

async def test_edgeql_with_alias_simple_03(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
C := (select Card { name_lower := str_lower(.name)})
FOR GROUP C
USING e := .element
BY e
IN g
UNION
{
key := e,
names := g.name_lower,
};
""",
tb.bag([
{'key': 'Water', 'names': ['bog monster', 'giant turtle']},
{'key': 'Fire', 'names': ['imp', 'dragon']},
{'key': 'Earth', 'names': ['dwarf', 'golem']},
{'key': 'Air', 'names': ['sprite', 'giant eagle', 'djinn']},
])
)

async def test_edgeql_with_alias_simple_04(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
C := (select Card { number := 9 })
FOR GROUP C { number }
USING e := .element
BY e
IN g
UNION
{
key := e,
numbers := g.number,
};
""",
tb.bag([
{'key': 'Water', 'numbers': {9, 9}},
{'key': 'Fire', 'numbers': {9, 9}},
{'key': 'Earth', 'numbers': {9, 9}},
{'key': 'Air', 'numbers': {9, 9, 9}},
])
)

async def test_edgeql_with_alias_for_01(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
U := (for x in {8, 9} select User { a := x }),
FOR GROUP U { name, a }
USING e := .name
BY e
IN g
UNION
{ e_ := e, g_ := g} {
key := .e_,
numbers := .g_.a,
};
""",
tb.bag([
{'key': 'Alice', 'numbers': {8, 9}},
{'key': 'Bob', 'numbers': {8, 9}},
{'key': 'Carol', 'numbers': {8, 9}},
{'key': 'Dave', 'numbers': {8, 9}},
])
)

async def test_edgeql_with_alias_for_02(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
X := {8, 9},
U := (for x in X select User { a := x }),
FOR GROUP U
USING e := .name
BY e
IN g
UNION
{ e_ := e, g_ := g} {
key := .e_,
numbers := .g_.a,
};
""",
tb.bag([
{'key': 'Alice', 'numbers': {8, 9}},
{'key': 'Bob', 'numbers': {8, 9}},
{'key': 'Carol', 'numbers': {8, 9}},
{'key': 'Dave', 'numbers': {8, 9}},
])
)

async def test_edgeql_with_alias_for_03(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
X := {8, 9},
FOR GROUP (for x in X select User { a := x })
USING e := .name
BY e
IN g
UNION
{ e_ := e, g_ := g} {
key := .e_,
numbers := .g_.a,
};
""",
tb.bag([
{'key': 'Alice', 'numbers': {8, 9}},
{'key': 'Bob', 'numbers': {8, 9}},
{'key': 'Carol', 'numbers': {8, 9}},
{'key': 'Dave', 'numbers': {8, 9}},
])
)

async def test_edgeql_with_alias_for_04(self):
await self.assert_query_result(
r"""
WITH
MODULE cards,
U := (
for x in enumerate({8, 9})
select User { a := x.0, b := x.1 }
),
FOR GROUP U
USING e := .name
BY e
IN g
UNION
{ e_ := e, g_ := g} {
key := .e_,
numbers := (.g_.a, .g_.b),
};
""",
tb.bag([
{'key': 'Alice', 'numbers': {(0, 8), (1, 9)}},
{'key': 'Bob', 'numbers': {(0, 8), (1, 9)}},
{'key': 'Carol', 'numbers': {(0, 8), (1, 9)}},
{'key': 'Dave', 'numbers': {(0, 8), (1, 9)}},
])
)

0 comments on commit c83c843

Please sign in to comment.