relationalai 0.13.2__py3-none-any.whl → 0.13.4__py3-none-any.whl

relationalai/semantics/lqp/model2lqp.py

@@ -1,22 +1,31 @@
 from relationalai.semantics.metamodel import ir, builtins, helpers, types
 from relationalai.semantics.metamodel.visitor import collect_by_type
-from relationalai.semantics.metamodel.util import FrozenOrderedSet
-from relationalai.semantics.lqp import ir as lqp, utils, types as lqp_types, builtins as lqp_builtins
+from relationalai.semantics.metamodel.util import FrozenOrderedSet, OrderedSet
+from relationalai.semantics.metamodel.compiler import group_tasks
+from relationalai.semantics.lqp import ir as lqp, utils, types as lqp_types
 from relationalai.semantics.lqp.primitives import lqp_avg_op, lqp_operator, build_primitive
 from relationalai.semantics.lqp.pragmas import pragma_to_lqp_name
 from relationalai.semantics.lqp.types import meta_type_to_lqp
 from relationalai.semantics.lqp.constructors import (
     mk_abstraction, mk_and, mk_exists, mk_or, mk_pragma, mk_primitive,
-    mk_specialized_value, mk_type, mk_value, mk_attribute
+    mk_specialized_value, mk_type, mk_value, mk_attribute,
 )
-from relationalai.semantics.lqp.utils import TranslationCtx, gen_unique_var
+from relationalai.semantics.lqp.algorithms import (
+    is_script, is_algorithm_script, is_algorithm_logical,
+    is_while_loop, is_while_script, construct_monoid, is_empty_instruction
+)
+from relationalai.semantics.lqp.builtins import (
+    has_global_annotation, get_upsert_annotation, get_monoid_annotation,
+    get_monus_annotation, has_assign_annotation, get_arity, supported_lqp_annotations
+)
+from relationalai.semantics.lqp.utils import TranslationCtx, ExportDescriptor, gen_unique_var, gen_rel_id
 from relationalai.semantics.lqp.validators import assert_valid_input
 from relationalai.semantics.lqp.rewrite.functional_dependencies import (
     normalized_fd, contains_only_declarable_constraints
 )
 from decimal import Decimal as PyDecimal
 from datetime import datetime, date, timezone
-from typing import Tuple, cast, Union, Optional
+from typing import Sequence, Tuple, cast, Union, Optional
 from warnings import warn
 import re
 import uuid
@@ -42,8 +51,8 @@ def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[
     reads.extend(_get_output_reads(ctx.output_ids))
 
     export_info = None
-    if len(ctx.export_ids) > 0:
-        export_filename, col_types, export_reads = _get_export_reads(ctx.export_ids)
+    if len(ctx.export_descriptors) > 0:
+        export_filename, col_types, export_reads = _get_export_reads(ctx.export_descriptors)
         reads.extend(export_reads)
         export_info = (export_filename, col_types)
 
@@ -76,14 +85,17 @@ def _get_output_reads(output_ids: list[tuple[lqp.RelationId, str]]) -> list[lqp.
         reads.append(lqp.Read(read_type=output, meta=None))
     return reads
 
-def _get_export_reads(export_ids: list[tuple[lqp.RelationId, int, lqp.Type]]) -> tuple[str, list, list[lqp.Read]]:
+def _get_export_reads(descriptors: list[ExportDescriptor]) -> tuple[str, list, list[lqp.Read]]:
     reads = []
     csv_columns = []
     col_info = []
-    for (col_id, col_num, col_type) in sorted(export_ids, key=lambda x: x[1]):
-        col_name = f"col{col_num:03}"
-        csv_columns.append(lqp.ExportCSVColumn(column_name=col_name, column_data=col_id, meta=None))
-        col_info.append((col_name, col_type))
+    for descriptor in sorted(descriptors, key=lambda x: x.column_number):
+        csv_columns.append(lqp.ExportCSVColumn(
+            column_name=descriptor.column_name,
+            column_data=descriptor.relation_id,
+            meta=None,
+        ))
+        col_info.append((descriptor.column_name, descriptor.column_type))
 
     # Generate a random name for the internal export path
     export_filename = "export_" + str(uuid.uuid4()).replace("-", "_")
@@ -106,6 +118,8 @@ def _get_export_reads(export_ids: list[tuple[lqp.RelationId, int, lqp.Type]]) ->
 def _translate_to_decls(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
     if contains_only_declarable_constraints(rule):
         return _translate_to_constraint_decls(ctx, rule)
+    elif is_algorithm_logical(rule):
+        return _translate_algorithms(ctx, rule)
     else:
         return _translate_to_standard_decl(ctx, rule)
 
@@ -128,8 +142,11 @@ def _translate_to_constraint_decls(ctx: TranslationCtx, rule: ir.Logical) -> lis
         lqp_guard = mk_abstraction(lqp_typed_vars, mk_and(lqp_guard_atoms))
         lqp_keys:list[lqp.Var] = [var for (var, _) in lqp_typed_keys] # type: ignore
         lqp_values:list[lqp.Var] = [var for (var, _) in lqp_typed_values] # type: ignore
+        lqp_id = utils.lqp_hash(fd.canonical_str)
+        lqp_name:lqp.RelationId = lqp.RelationId(id=lqp_id, meta=None)
 
         fd_decl = lqp.FunctionalDependency(
+            name=lqp_name,
             guard=lqp_guard,
             keys=lqp_keys,
             values=lqp_values,
@@ -140,39 +157,161 @@ def _translate_to_constraint_decls(ctx: TranslationCtx, rule: ir.Logical) -> lis
 
     return constraint_decls
 
-def _translate_to_standard_decl(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
-    effects = collect_by_type((ir.Output, ir.Update), rule)
-    aggregates = collect_by_type(ir.Aggregate, rule)
-    ranks = collect_by_type(ir.Rank, rule)
+def _translate_algorithms(ctx: TranslationCtx, task: ir.Logical) -> list[lqp.Declaration]:
+    assert is_algorithm_logical(task)
+    decls: list[lqp.Declaration] = []
+    for subtask in task.body:
+        assert is_algorithm_script(subtask), "Expected all subtasks to be algorithm scripts"
+        decls.extend(_translate_algorithm_script(ctx, subtask))
+    return decls
+
+def _translate_algorithm_script(ctx: TranslationCtx, alg_task: ir.Sequence) -> list[lqp.Declaration]:
+    assert is_algorithm_script(alg_task), "Expected Sequence @algorithm @script "
+
+    alg_globals = _find_algorithm_global_relation_ids(ctx, alg_task)
+    alg_body = _translate_script(ctx, alg_task)
+
+    return [lqp.Algorithm(global_=alg_globals, body=alg_body, meta=None)]
+
+def _find_algorithm_global_relation_ids(ctx: TranslationCtx, alg_task: ir.Sequence) -> list[lqp.RelationId]:
+    result = []
+    updates = collect_by_type(ir.Update, alg_task)
+    for update in updates:
+        if has_global_annotation(update):
+            bindings = _effect_bindings(update)
+            projection, _ = _translate_bindings(ctx, bindings)
+            rel_id = get_relation_id(ctx, update.relation, projection)
+            result.append(rel_id)
+    return list(dict.fromkeys(result))
+
+def _translate_script(ctx: TranslationCtx, task: ir.Sequence) -> lqp.Script:
+    assert is_script(task), "Expected a @script Sequence"
+
+    constructs: list[lqp.Construct] = []
+
+    for subtask in task.tasks:
+        if is_empty_instruction(subtask):
+            constructs.append(_translate_empty_instruction(ctx, subtask))
+        elif isinstance(subtask, ir.Logical):
+            constructs.extend(_translate_instruction(ctx, subtask))
+        elif isinstance(subtask, ir.Break):
+            constructs.append(_translate_break_instruction(ctx, subtask))
+        elif is_while_loop(subtask):
+            constructs.append(_translate_while_loop(ctx, subtask))
+        else:
+            raise Exception(f"Unsupported script instruction: {subtask}")
+
+    return lqp.Script(constructs=constructs, meta=None)
+
+def _translate_while_loop(ctx: TranslationCtx, task: ir.Loop) -> lqp.Loop:
+    assert is_while_loop(task), "Expected a @while Loop"
+    assert len(task.iter) == 0, "Temporalized loops not supported"
+
+    while_script_task = task.body
+    assert is_while_script(while_script_task), "The body of a @while Loop must be a @while @script Sequence"
+    body_script = _translate_script(ctx, while_script_task)
+
+    # No init instructions in the translation of PyRel Loops to to LQP loops
+    return lqp.Loop(init=[], body=body_script, meta=None)
+
+def _translate_break_instruction(ctx: TranslationCtx, task: ir.Break) -> lqp.Construct:
+    body = _translate_to_formula(ctx, task.check)
+
+    ctx.break_rule_counter += 1
+
+    rel_id = gen_rel_id(ctx, "break_cond_" + str(ctx.break_rule_counter))
+    return lqp.Break(
+        name = rel_id,
+        body = mk_abstraction([], body),
+        attrs = [],
+        meta = None,
+    )
+
+def _translate_empty_instruction(ctx: TranslationCtx, rule: ir.Logical) -> lqp.Instruction:
+    assert is_empty_instruction(rule)
+    updates = collect_by_type(ir.Update, rule)
+    assert len(updates) == 1
+    update = updates[0]
+    bindings = _effect_bindings(update)
+
+    # We need  to make sure that variable names have a leading underscore
+    normalized_bindings:Sequence[ir.Var] = []
+    for v in bindings:
+        assert isinstance(v, ir.Var)
+        if not v.name.startswith("_"):
+            v = ir.Var(v.type, "_" + v.name)
+        normalized_bindings.append(v)
+
+    projection, eqs = _translate_bindings(ctx, normalized_bindings)
+    assert len(eqs) == 0
+    rel_id = get_relation_id(ctx, update.relation, projection)
+    abstraction = mk_abstraction(projection, mk_or([])) # empty body = false
+    return lqp.Assign(name = rel_id, body = abstraction, attrs = [], meta = None)
+
+def _translate_instruction(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Instruction]:
+    effects = collect_by_type((ir.Update, ir.Output), rule)
 
-    # TODO: should this ever actually come in as input?
     if len(effects) == 0:
         return []
 
-    assert len(ranks) == 0 or len(aggregates) == 0, "rules cannot have both aggregates and ranks"
+    conjuncts = _translate_to_formula(ctx, rule)
+    res = []
+    for effect in effects:
+        assert isinstance(effect, ir.Update), f"Got an effect of type {type(effect)} in a loop, which is invalid."
 
-    conjuncts = []
-    for task in rule.body:
-        if isinstance(task, (ir.Output, ir.Update)):
-            continue
-        conjuncts.append(_translate_to_formula(ctx, task))
+        bindings = _effect_bindings(effect)
+        projection, eqs = _translate_bindings(ctx, bindings)
+
+        eqs.append(conjuncts)
+        new_body = mk_and(eqs)
+
+        rel_id = get_relation_id(ctx, effect.relation, projection)
+        abstraction = mk_abstraction(projection, new_body)
+
+        upsert = get_upsert_annotation(effect)
+        monoid = get_monoid_annotation(effect)
+        monus = get_monus_annotation(effect)
+
+        if has_assign_annotation(effect):
+            res.append(lqp.Assign(name = rel_id, body = abstraction, attrs = [], meta = None))
+        elif upsert is not None:
+            res.append(lqp.Upsert(value_arity=get_arity(upsert), name = rel_id, body = abstraction, attrs = [], meta = None))
+        elif monoid is not None:
+            res.append(lqp.MonoidDef(
+                value_arity=get_arity(monoid),
+                monoid=construct_monoid(monoid),
+                name = rel_id,
+                body = abstraction,
+                attrs = [],
+                meta = None
+            ))
+        elif monus is not None:
+            res.append(lqp.MonusDef(
+                value_arity=get_arity(monus),
+                monoid=construct_monoid(monus),
+                name = rel_id,
+                body = abstraction,
+                attrs = [],
+                meta = None
+            ))
+
+    return res
 
-    # Aggregates reduce over the body
-    if aggregates or ranks:
-        aggr_body = mk_and(conjuncts)
-        conjuncts = []
-        for aggr in aggregates:
-            conjuncts.append(_translate_aggregate(ctx, aggr, aggr_body))
-        for rank in ranks:
-            conjuncts.append(_translate_rank(ctx, rank, aggr_body))
+def _translate_to_standard_decl(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
+    effects = collect_by_type((ir.Output, ir.Update), rule)
 
-    return [_translate_effect(ctx, effect, mk_and(conjuncts)) for effect in effects]
+    # TODO: should this ever actually come in as input?
+    if len(effects) == 0:
+        return []
+
+    conjuncts = _translate_to_formula(ctx, rule)
+    return [_translate_effect(ctx, effect, conjuncts) for effect in effects]
 
 def _translate_annotations(annotations: FrozenOrderedSet[ir.Annotation]) -> list[lqp.Attribute]:
     attributes = []
     for annotation in annotations:
 
-        if annotation.relation.name in lqp_builtins.annotations_to_emit:
+        if annotation.relation.name in supported_lqp_annotations:
             if any(not isinstance(a, ir.Literal) for a in annotation.args):
                 warn("LQP currently ignores annotation parameters with non-literal values")
                 continue
@@ -191,14 +330,11 @@ def _translate_annotations(annotations: FrozenOrderedSet[ir.Annotation]) -> list
 def _translate_effect(ctx: TranslationCtx, effect: Union[ir.Output, ir.Update], body: lqp.Formula) -> lqp.Declaration:
     bindings = _effect_bindings(effect)
 
-    def _is_export(e):
-        return isinstance(e, ir.Output) and builtins.export_annotation in e.annotations
-
     if isinstance(effect, ir.Output):
         projection, eqs, suffix = _translate_output_bindings(ctx, bindings)
         meta_id = effect.id
 
-        if _is_export(effect):
+        if helpers.is_export(effect):
             def_name = "export_relation" + suffix
         else:
             def_name = "output" + suffix
@@ -214,14 +350,23 @@ def _translate_effect(ctx: TranslationCtx, effect: Union[ir.Output, ir.Update],
     new_body = mk_and(eqs)
 
     # Context bookkeeping for exports and outputs
-    if _is_export(effect):
+    if helpers.is_export(effect):
         # The row id is the first n-1 elements, and the actual data is the last element. Its
         # type is stored in the first element of the tuple.
         col_type = projection[-1][1]
         _col_num_match = re.search(r"export_relation_col([0-9]+)", def_name)
         assert _col_num_match, f"Could not find column number in suffix: {def_name}"
         col_num = int(_col_num_match.group(1))
-        ctx.export_ids.append((rel_id, col_num, col_type))
+        col_name = f"col{col_num}"
+        if isinstance(effect, ir.Output) and len(effect.aliases) > 0:
+            aliases_list = list(effect.aliases)
+            col_name = aliases_list[-1][0]
+        ctx.export_descriptors.append(ExportDescriptor(
+            relation_id=rel_id,
+            column_name=col_name,
+            column_number=col_num,
+            column_type=col_type,
+        ))
     elif isinstance(effect, ir.Output):
         ctx.output_ids.append((rel_id, def_name))
 
@@ -544,7 +689,36 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
 
 def _translate_to_formula(ctx: TranslationCtx, task: ir.Task) -> lqp.Formula:
     if isinstance(task, ir.Logical):
-        conjuncts = [_translate_to_formula(ctx, child) for child in task.body]
+        # For aggregates and ranks, the expected format is:
+        #
+        # Logical
+        #   body_task1
+        #   body_task2
+        #   ...
+        #   aggregate/rank task
+        #
+        # If we see that the Logical is in this format, it should be translated as an
+        # aggregate/rank node.
+        groups = group_tasks(task.body, {
+            "aggregates": ir.Aggregate,
+            "ranks": ir.Rank,
+        })
+
+        aggregates = groups.get("aggregates", OrderedSet[ir.Task]())
+        ranks = groups.get("ranks", OrderedSet[ir.Task]())
+
+        if aggregates or ranks:
+            conjuncts = []
+            body = mk_and([_translate_to_formula(ctx, t) for t in task.body])
+            for aggr in aggregates:
+                assert(isinstance(aggr, ir.Aggregate))
+                conjuncts.append(_translate_aggregate(ctx, aggr, body))
+            for rank in ranks:
+                assert(isinstance(rank, ir.Rank))
+                conjuncts.append(_translate_rank(ctx, rank, body))
+        else:
+            # If there are no aggregates or ranks, translate as a normal conjunction
+            conjuncts = [_translate_to_formula(ctx, child) for child in task.body]
         return mk_and(conjuncts)
     elif isinstance(task, ir.Lookup):
         return _translate_to_atom(ctx, task)
@@ -564,14 +738,13 @@ def _translate_to_formula(ctx: TranslationCtx, task: ir.Task) -> lqp.Formula:
 
         return mk_primitive("rel_primitive_hash_tuple_uint128", [v for v, _ in terms])
     elif isinstance(task, ir.Union):
-        # TODO: handle hoisted vars if needed
         disjs = [_translate_to_formula(ctx, child) for child in task.tasks]
         return mk_or(disjs)
-    elif isinstance(task, (ir.Aggregate, ir.Output, ir.Update)):
+    elif isinstance(task, (ir.Output, ir.Update)):
         # Nothing to do here, handled in _translate_to_decls
         return mk_and([])
-    elif isinstance(task, ir.Rank):
-        # Nothing to do here, handled in _translate_to_decls
+    elif isinstance(task, (ir.Aggregate, ir.Rank)):
+        # Nothing to do here, handled at the Logical level
         return mk_and([])
     else:
         raise NotImplementedError(f"Unknown task type (formula): {type(task)}")
@@ -666,7 +839,7 @@ def get_output_id(ctx: TranslationCtx, orig_name: str, metamodel_id: int) -> lqp
     unique_name = ctx.output_names.get_name_by_id(metamodel_id, orig_name)
     return utils.gen_rel_id(ctx, unique_name)
 
-def _translate_bindings(ctx: TranslationCtx, bindings: list[ir.Value]) -> Tuple[list[Tuple[lqp.Var, lqp.Type]], list[lqp.Formula]]:
+def _translate_bindings(ctx: TranslationCtx, bindings: Sequence[ir.Value]) -> Tuple[list[Tuple[lqp.Var, lqp.Type]], list[lqp.Formula]]:
     lqp_vars = []
     conjuncts = []
     for binding in bindings: