Replace ExprOpAggregate with Op with sig set

Author: andreypopp

qc0/compile.py

@@ -19,7 +19,6 @@
     RelAroundParent,
     Expr,
     ExprOp,
-    ExprOpAggregate,
     ExprRecord,
     ExprColumn,
     ExprIdentity,
@@ -264,8 +263,8 @@ def build_kernel():
         at = from_obj.at
         columns = []
         for field in rel.fields.values():
-            expr, from_obj = expr_to_sql(
-                field.expr,
+            expr, from_obj = op_to_sql(
+                field.op,
                 from_obj=from_obj.replace(at=at),
             )
             columns.append(expr.label(field.name))
@@ -292,11 +291,17 @@ def build_kernel():
         return None, from_obj
 
     result_columns = [from_obj.current.columns[c.name] for c in tuple(columns)]
-    for name, expr in rel.aggregates.items():
-        assert isinstance(expr, ExprOpAggregate), str(expr)
+    for name, op in rel.aggregates.items():
+        assert op.sig is not None
         columns, kernel = build_kernel()
-        value, inner_from_obj = op_to_sql(expr.op, from_obj=kernel)
-        value = expr.sig.compile([value])
+
+        value, inner_from_obj = rel_to_sql(op.rel, from_obj=kernel)
+        if op.expr is not None:
+            value, inner_from_obj = expr_to_sql(
+                op.expr, from_obj=inner_from_obj
+            )
+
+        value = op.sig.compile([value])
         if kernel.current in inner_from_obj.current._from_objects:
             cols = columns
         else:
@@ -312,7 +317,7 @@ def build_kernel():
             inner_sel, *((c.name, c.name) for c in columns), outer=True
         )
         result_columns.append(
-            sa.func.coalesce(inner_at.c.value, expr.sig.unit).label(name)
+            sa.func.coalesce(inner_at.c.value, op.sig.unit).label(name)
         )
 
     from_obj = From.make(
@@ -322,29 +327,14 @@ def build_kernel():
 
 
 def op_to_sql(op: Op, from_obj):
-    expr, from_obj = rel_to_sql(op.rel, from_obj=from_obj)
+    expr, inner_from_obj = rel_to_sql(op.rel, from_obj=from_obj)
     if op.expr is not None:
-        expr, from_obj = expr_to_sql(op.expr, from_obj=from_obj)
-    return expr, from_obj
-
-
-@singledispatch
-def expr_to_sql(expr: Expr, from_obj):
-    raise NotImplementedError(  # pragma: no cover
-        f"expr_to_sql({type(expr).__name__})"
-    )
-
-
-@expr_to_sql.register
-def ExprOp_to_sql(expr: ExprOp, from_obj):
-    return op_to_sql(expr.op, from_obj)
+        expr, inner_from_obj = expr_to_sql(op.expr, from_obj=inner_from_obj)
 
+    if op.sig is None:
+        return expr, inner_from_obj
 
-@expr_to_sql.register
-def ExprOpAggregate_to_sql(op: ExprOpAggregate, from_obj):
-    expr, inner_from_obj = op_to_sql(op.op, from_obj)
     if inner_from_obj.limit is not None or inner_from_obj.order is not None:
-        expr, inner_from_obj = op_to_sql(op.op, from_obj)
         inner_from_obj = From.make(inner_from_obj.to_select(expr).alias())
         value = sa.func.coalesce(
             op.sig.compile([inner_from_obj.at.c.value]),
@@ -363,15 +353,25 @@ def ExprOpAggregate_to_sql(op: ExprOpAggregate, from_obj):
     return at.c.value, from_obj
 
 
+@singledispatch
+def expr_to_sql(expr: Expr, from_obj):
+    raise NotImplementedError(  # pragma: no cover
+        f"expr_to_sql({type(expr).__name__})"
+    )
+
+
+@expr_to_sql.register
+def ExprOp_to_sql(expr: ExprOp, from_obj):
+    return op_to_sql(expr.op, from_obj)
+
+
 @expr_to_sql.register
 def ExprRecord_to_sql(op: ExprRecord, from_obj):
     args = []
     at = from_obj.at
     for field in op.fields.values():
         args.append(sa.literal(field.name))
-        expr, from_obj = expr_to_sql(
-            field.expr, from_obj=from_obj.replace(at=at)
-        )
+        expr, from_obj = op_to_sql(field.op, from_obj=from_obj.replace(at=at))
         args.append(expr)
     return sa.func.jsonb_build_object(*args), from_obj
 

qc0/op.py

@@ -14,7 +14,7 @@
 from sqlalchemy import Table, ForeignKey
 from sqlalchemy.sql.elements import ColumnClause
 from .base import Struct, undefined
-from .scope import Scope, Cardinality
+from .scope import Scope, EmptyScope, Cardinality
 from .syntax import Syn, Compose
 
 
@@ -24,6 +24,7 @@ class Op(Struct):
     rel: Rel
     expr: Optional[Expr]
     syn: Optional[Syn] = None
+    sig: Any = None
 
     def grow_expr(self, expr, syn=None, scope=undefined, card=undefined):
         card = self.card if card is undefined else card
@@ -49,6 +50,9 @@ def grow_rel(self, rel, syn=None, scope=undefined, card=undefined):
         )
         return self.replace(rel=rel, scope=scope, card=card, syn=syn)
 
+    def aggregate(self, sig):
+        return self.replace(sig=sig, scope=EmptyScope(), card=Cardinality.ONE)
+
     def __yaml__(self):
         rep = super(Op, self).__yaml__()
         rep.pop("scope")
@@ -119,11 +123,6 @@ class ExprOp(Expr):
     op: Op
 
 
-class ExprOpAggregate(Expr):
-    op: Op
-    sig: Any
-
-
 class ExprRecord(Expr):
     fields: Dict[str, Field]
 
@@ -149,7 +148,7 @@ class ExprApply(Expr):
 
 class Field(Struct):
     name: str
-    expr: Expr
+    op: Op
 
 
 class Sort(Struct):

qc0/plan.py

@@ -52,7 +52,6 @@
     RelAroundParent,
     Expr,
     ExprOp,
-    ExprOpAggregate,
     ExprRecord,
     ExprColumn,
     ExprIdentity,
@@ -99,12 +98,10 @@ def build_op_expr(op: Op):
         if isinstance(op.scope, RecordScope):
             fields = {}
             for name, f in op.scope.fields.items():
-                fop = build_op(f.syn, make_parent(op.scope.parent))
-                if fop.card == Cardinality.SEQ:
-                    expr = ExprOpAggregate(op=fop, sig=JsonAggSig)
-                else:
-                    expr = ExprOp(op=fop)
-                fields[name] = Field(expr=expr, name=name)
+                field_op = build_op(f.syn, make_parent(op.scope.parent))
+                if field_op.card == Cardinality.SEQ:
+                    field_op = field_op.aggregate(JsonAggSig)
+                fields[name] = Field(op=field_op, name=name)
             return Op(
                 rel=op.rel,
                 expr=ExprRecord(fields=fields),
@@ -135,14 +132,12 @@ def build_op_expr(op: Op):
             for name, f in op.scope.fields.items():
                 fields[name] = Field(
                     name=name,
-                    expr=ExprOp(
-                        Op(
-                            expr=ExprColumn(column=sa.column(name)),
-                            rel=RelParent(parent=op),
-                            card=Cardinality.ONE,
-                            scope=EmptyScope(),
-                            syn=Nav(name),
-                        )
+                    op=Op(
+                        expr=ExprColumn(column=sa.column(name)),
+                        rel=RelParent(parent=op),
+                        card=Cardinality.ONE,
+                        scope=EmptyScope(),
+                        syn=Nav(name),
                     ),
                 )
             expr = ExprRecord(fields=fields)
@@ -153,7 +148,6 @@ def build_op_expr(op: Op):
                 scope=op.scope,
                 syn=None,
             )
-        assert False, f"unable to build an expr at this scope: {op!r}"
     return op
 
 
@@ -262,12 +256,10 @@ def Nav_to_op(syn: Nav, parent: Op):
 
             def wrap(op):
                 if op.card == Cardinality.SEQ:
-                    expr = ExprOpAggregate(op=op, sig=JsonAggSig)
+                    op = op.aggregate(JsonAggSig)
                 else:
-                    assert isinstance(op.rel, RelVoid)
-                    assert isinstance(op.expr, ExprOpAggregate)
-                    expr = op.expr
-                name = parent.scope.add_aggregate(expr)
+                    assert op.sig is not None
+                name = parent.scope.add_aggregate(op)
                 return parent.grow_expr(
                     expr=ExprColumn(column=sa.column(name)),
                     scope=EmptyScope(),
@@ -391,9 +383,7 @@ def Apply_to_op(syn: Apply, parent: Op):
             if op.expr is None:
                 if isinstance(op.scope, TableScope):
                     op = op.grow_expr(ExprIdentity(table=op.scope.table))
-                else:
-                    assert False, f"{syn.name}: unable to group by this"
-            fields[name] = Field(expr=ExprOp(op), name=name)
+            fields[name] = Field(op=op, name=name)
 
         rel = RelGroup(
             rel=parent.rel,
@@ -407,13 +397,7 @@ def Apply_to_op(syn: Apply, parent: Op):
         sig = AggrSig.get(syn.name)
         if sig:
             assert parent.card >= Cardinality.SEQ, parent
-            return Op(
-                expr=ExprOpAggregate(op=parent, sig=sig),
-                rel=RelVoid(),
-                scope=EmptyScope(),
-                card=Cardinality.ONE,
-                syn=syn,
-            )
+            return parent.aggregate(sig)
 
         sig = FuncSig.get(syn.name)
         if sig:

qc0/scope.py

@@ -96,10 +96,10 @@ class GroupScope(Scope):
     fields: Dict[str, Field]
     aggregates: Dict[str, Any]
 
-    def add_aggregate(self, expr):
+    def add_aggregate(self, op):
         idx = len(self.aggregates)
         name = f"aggr_{idx}"
-        self.aggregates[name] = expr
+        self.aggregates[name] = op
         return name
 
     def __yaml__(self):