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

relationalai/semantics/metamodel/rewrite/flatten.py

@@ -13,7 +13,8 @@ class Flatten(Pass):
     """
     Traverses the model's root to flatten it as much as possible. The result of this pass is
     a Logical root where all nested tasks that represent a rule in Rel are extracted to the
-    top level.
+    top level. Additionally, any Sequence is promoted to the top level Logical (but
+    encapsulated by a Logical).
 
     - nested logical with updates becomes a top-level logical (a rule)
 
@@ -122,6 +123,35 @@ class Flatten(Pass):
             Logical
                 lookup tmp2
                 output
+
+    - a Sequence is promoted to the top level Logical, encapsulated by a Logical:
+    From:
+        Logical
+            Logical
+                lookup
+                derive foo
+            Sequence
+                Logical
+                    ...
+                Loop
+                    Sequence
+                        ...
+                Logical
+                    ...
+    To:
+        Logical
+            Logical
+                lookup
+                derive foo
+            Logical
+                Sequence
+                    Logical
+                        ...
+                    Loop
+                        Sequence
+                            ...
+                    Logical
+                        ...
     """
 
     def __init__(self, use_sql: bool=False):
@@ -181,11 +211,8 @@ class Flatten(Pass):
     def handle(self, task: ir.Task, ctx: Context) -> Flatten.HandleResult:
         if isinstance(task, ir.Logical):
             return self.handle_logical(task, ctx)
-        elif isinstance(task, ir.Union) and (task.hoisted or self._use_sql):
-            # Only flatten Unions which hoist variables. If there are no hoisted variables,
-            # then the Union acts as a filter, and it can be inefficient to flatten it.
-            #
-            # However, for the SQL backend, we always need to flatten Unions for correct SQL
+        elif isinstance(task, ir.Union) and self._use_sql:
+            # The SQL backend needs to flatten Unions for correct SQL
             # generation.
             return self.handle_union(task, ctx)
         elif isinstance(task, ir.Match):
@@ -194,6 +221,8 @@ class Flatten(Pass):
             return self.handle_require(task, ctx)
         elif isinstance(task, ir.Not):
             return self.handle_not(task, ctx)
+        elif isinstance(task, ir.Sequence):
+            return self.handle_sequence(task, ctx)
         else:
             return Flatten.HandleResult(task)
 
@@ -253,9 +282,9 @@ class Flatten(Pass):
 
             for output in groups["outputs"]:
                 assert(isinstance(output, ir.Output))
-                new_body = info.task_dependencies(output)
-                new_body.update(ctx.extra_tasks)
-                new_body.add(output)
+                new_body = OrderedSet.from_iterable(t.clone() for t in info.task_dependencies(output))
+                new_body.update(t.clone() for t in ctx.extra_tasks)
+                new_body.add(output.clone())
                 ctx.rewrite_ctx.top_level.append(ir.Logical(task.engine, task.hoisted, tuple(new_body), task.annotations))
 
             return Flatten.HandleResult(None)
@@ -263,9 +292,9 @@ class Flatten(Pass):
         # if there are updates, extract as a new top level rule
         if groups["updates"]:
             # add task dependencies to the body
-            body.prefix(ctx.info.task_dependencies(task))
+            body.prefix(t.clone() for t in ctx.info.task_dependencies(task))
             # potentially add context extra tasks
-            body.update(ctx.extra_tasks)
+            body.update(t.clone() for t in ctx.extra_tasks)
             ctx.rewrite_ctx.top_level.append(ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations))
             return Flatten.HandleResult(None)
 
@@ -278,7 +307,7 @@ class Flatten(Pass):
             agg = cast(ir.Aggregate, groups["aggregates"].some())
 
             # add agg dependencies to the body
-            body.prefix(ctx.info.task_dependencies(agg))
+            body.prefix(t.clone() for t in ctx.info.task_dependencies(agg))
 
             # extract a new logical for the aggregate, exposing aggregate group-by and results
             exposed_vars = OrderedSet.from_iterable(list(agg.group) + helpers.aggregate_outputs(agg))
@@ -298,7 +327,7 @@ class Flatten(Pass):
             rank = cast(ir.Rank, groups["ranks"].some())
 
             # add rank dependencies to the body
-            body.prefix(ctx.info.task_dependencies(rank))
+            body.prefix(t.clone() for t in ctx.info.task_dependencies(rank))
             # for rank, we sort by the args, but the result includes the keys to preserve bag semantics.
             exposed_vars_raw = list(rank.projection) + list(rank.group) + list(rank.args) +[rank.result]
             # deduplicate vars
@@ -487,6 +516,14 @@ class Flatten(Pass):
                 task.annotations
             ))
 
+    def handle_sequence(self, task: ir.Sequence, ctx: Context):
+        new_logical = f.logical(
+            body = [task],
+            engine = task.engine
+        )
+        ctx.rewrite_ctx.top_level.append(new_logical)
+        return Flatten.HandleResult(None)
+
 #--------------------------------------------------
 # Helpers
 #--------------------------------------------------

relationalai/semantics/metamodel/rewrite/format_outputs.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from typing import Tuple
 
-from relationalai.semantics.metamodel import builtins, ir, factory as f, types, visitor
+from relationalai.semantics.metamodel import builtins, ir, factory as f, types, visitor, helpers
 from relationalai.semantics.metamodel.compiler import Pass, group_tasks
 from relationalai.semantics.metamodel.util import OrderedSet
 from relationalai.semantics.metamodel.util import FrozenOrderedSet
@@ -63,7 +63,7 @@ def adjust_outputs(task: ir.Logical, outputs: OrderedSet[ir.Task], wide_outputs:
                 # Remove the original output. This is replaced by per-column outputs below
                 body.remove(output)
 
-                is_export = builtins.export_annotation in output.annotations
+                is_export = helpers.is_export(output)
 
                 # Generate an output for each "column"
                 # output looks like def output(:cols, :col000, key0, key1, value):
@@ -100,7 +100,13 @@ def _generate_output_column(output: ir.Output, idx: int, alias: tuple[str, ir.Va
         (not is_primitive(alias[1].type) or alias[1].type == types.Hash)):
 
         uuid = f.var(f"{alias[0]}_{idx}_uuid", types.String)
-        aliases.append((uuid.name, uuid))
+
+        if not is_primitive(alias[1].type):
+            # For non-primitive types, we keep the original alias
+            aliases.append((alias[0], uuid))
+        else:
+            # For Hash types, we use the uuid name as alias
+            aliases.append((uuid.name, uuid))
 
         return [
             ir.Lookup(None, builtins.uuid_to_string, (alias[1], uuid)),

relationalai/semantics/metamodel/typer/checker.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass, field
 from typing import Optional, List, Union as PyUnion, Tuple, cast
 
 from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
-from relationalai.semantics.metamodel import ir, types, visitor, compiler
+from relationalai.semantics.metamodel import ir, types, visitor, compiler, executor
 import rich
 
 
@@ -39,7 +39,8 @@ class CheckEnv:
 
     def _complain(self, node: ir.Node, msg: str):
         """Report an error."""
-        self.diags.append(CheckError(msg, node))
+        if not executor.SUPPRESS_TYPE_ERRORS:
+            self.diags.append(CheckError(msg, node))
 
 
 @dataclass
@@ -306,8 +307,9 @@ class CheckModel(visitor.DAGVisitor):
         for x in node.hoisted:
             if not CheckModel._variable_occurs_in(x, node.body):
                 self.env._complain(node, f"Variable {ir.node_to_string(x).strip()} is hoisted but not used in the body of {ir.node_to_string(node).strip()}.")
-        if not CheckModel._variable_occurs_in(node.iter, node.body):
-            self.env._complain(node, f"Variable {node.iter} is the loop iterator but is not used in the body of {ir.node_to_string(node).strip()}.")
+        for iter_var in node.iter:
+            if not CheckModel._variable_occurs_in(iter_var, node.body):
+                self.env._complain(node, f"Variable {iter_var} is the loop iterator but is not used in the body of {ir.node_to_string(node).strip()}.")
         return super().visit_loop(node, parent)
 
     def visit_update(self, node: ir.Update, parent: Optional[ir.Node]=None):

relationalai/semantics/metamodel/typer/typer.py

@@ -6,7 +6,7 @@ import datetime
 from decimal import Decimal as PyDecimal
 from typing import Optional, Union, Tuple
 from relationalai import debugging
-from relationalai.semantics.metamodel import builtins, helpers, ir, types, visitor, compiler, factory as f
+from relationalai.semantics.metamodel import builtins, helpers, ir, types, visitor, compiler, factory as f, executor
 from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
 import rich
 import sys
@@ -1361,9 +1361,6 @@ class Replacer(visitor.Rewriter):
 # Typer pass
 #--------------------------------------------------
 
-# global flag to suppress type errors from being printed
-SUPPRESS_TYPE_ERRORS = False
-
 class InferTypes(compiler.Pass):
     def __init__(self):
         super().__init__()
@@ -1392,7 +1389,7 @@ class InferTypes(compiler.Pass):
         with debugging.span("type.replace"):
             final = Replacer(w.net).walk(model)
 
-        if not SUPPRESS_TYPE_ERRORS:
+        if not executor.SUPPRESS_TYPE_ERRORS:
             for err in w.net.errors:
                 rich.print(str(err), file=sys.stderr)
 

relationalai/semantics/metamodel/visitor.py

@@ -466,7 +466,8 @@ class Visitor(GenericVisitor[None]):
             self._walk_engine(node.engine, node)
         for h in node.hoisted:
             self._walk_var_or_default(h, node)
-        self._walk_var(node.iter, node)
+        for iter in node.iter:
+            self._walk_var(iter, node)
         self._walk_node(node.body, node)
         for a in node.annotations:
             self._walk_node(a, node)
@@ -935,9 +936,9 @@ class Rewriter():
     #
     def handle_loop(self, node: ir.Loop, parent: ir.Node):
         hoisted = rewrite_list(ir.VarOrDefault, lambda n: self.walk(n, node), node.hoisted)
-        iter_val = self.walk(node.iter, node)
+        iter = rewrite_list(ir.Var, lambda n: self.walk(n, node), node.iter)
         body = self.walk(node.body, node)
-        return node.reconstruct(node.engine, hoisted, iter_val, body, node.annotations)
+        return node.reconstruct(node.engine, hoisted, iter, body, node.concurrency, node.annotations)
 
     def handle_break(self, node: ir.Break, parent: ir.Node):
         check = self.walk(node.check, node)

relationalai/semantics/reasoners/optimization/solvers_dev.py

@@ -333,7 +333,7 @@ class SolverModelDev:
         executor.execute_raw(textwrap.dedent(f"""
         def delete[:{self.point._name}]: {self.point._name}
         def insert(:{self.point._name}, var, val): {self.points._name}(int128[{i}], var, val)
-        """), readonly=False)
+        """))
         return None
 
     # print summary of the solver result

relationalai/semantics/reasoners/optimization/solvers_pb.py

@@ -599,7 +599,7 @@ class SolverModelPB:
         }}]}}
         """)
 
-        executor.execute_raw(export_rel, readonly=False, query_timeout_mins=query_timeout_mins)
+        executor.execute_raw(export_rel, query_timeout_mins=query_timeout_mins)
 
     def _import_solver_results_from_csv(
         self,
@@ -695,7 +695,7 @@ class SolverModelPB:
         }}
         """)
 
-        executor.execute_raw(load_and_extract_rel, readonly=False, query_timeout_mins=query_timeout_mins)
+        executor.execute_raw(load_and_extract_rel, query_timeout_mins=query_timeout_mins)
 
     def _export_model_to_protobuf(
         self,
@@ -791,7 +791,6 @@ class SolverModelPB:
 
         executor.execute_raw(
             textwrap.dedent(extract_rel) + textwrap.dedent(insert_points_relation),
-            readonly=False,
             query_timeout_mins=query_timeout_mins
         )
 
@@ -929,7 +928,7 @@ class SolverModelPB:
         def delete[:{self.point._name}]: {self.point._name}
         def insert(:{self.point._name}, variable, value): {self.points._name}(int128[{point_index}], variable, value)
         """
-        executor.execute_raw(textwrap.dedent(load_point_relation), readonly=False)
+        executor.execute_raw(textwrap.dedent(load_point_relation))
 
     def summarize_result(self) -> Any:
         """Print solver result summary.

relationalai/semantics/rel/compiler.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+
 from typing import Any, Iterable, Sequence as PySequence, cast, Tuple, Union
 from dataclasses import dataclass, field
 from decimal import Decimal as PyDecimal
@@ -843,7 +844,7 @@ class ModelToRel:
 
     def _effect_name(self, n: ir.Task):
         """ Return the name to be used for the effect (e.g. the relation name, output, etc). """
-        if isinstance(n, ir.Output) and bt.export_annotation in n.annotations:
+        if helpers.is_export(n):
             return "Export_Relation"
         elif isinstance(n, ir.Output):
             return "output"

relationalai/semantics/rel/executor.py

@@ -355,8 +355,9 @@ class RelExecutor(e.Executor):
 
     # NOTE(coey): this is added temporarily to support executing Rel for the solvers library in EA.
     # It can be removed once this is no longer needed by the solvers library.
-    def execute_raw(self, raw_rel:str, readonly:bool=True, query_timeout_mins:int|None=None) -> DataFrame:
-        raw_results = self.resources.exec_raw(self.database, self.engine, raw_rel, readonly, nowait_durable=True, query_timeout_mins=query_timeout_mins)
+    def execute_raw(self, raw_rel:str, query_timeout_mins:int|None=None) -> DataFrame:
+        # NOTE intentionally hard-coding to read-only=False, because read-only Rel queries are deprecated.
+        raw_results = self.resources.exec_raw(self.database, self.engine, raw_rel, False, nowait_durable=True, query_timeout_mins=query_timeout_mins)
         df, errs = result_helpers.format_results(raw_results, None, generation=Generation.QB)  # Pass None for task parameter
         self.report_errors(errs)
         return df

relationalai/semantics/tests/lqp/algorithms.py

@@ -0,0 +1,345 @@
+"""
+Constructing Metamodel IR with Algorithms
+
+We introduce a set of programmatic constructs that provide a convenient syntax for
+constructing PyRel's metamodel IR representations for Loopy algorithms. Importantly, these
+macros construct a new model using PyRel declarations constructed with a _base model_. The
+base model needs to be also used to declare all concepts and relationships.
+
+Below we illustrate the use of these macros by constructing a simple reachability
+algorithm, whose Rel-like pseudo-code is as follows:
+
+```
+algorithm
+    setup
+        def edge = { (1,2); (2,3); (3,4) }
+        def source = { 1 }
+    end setup
+    @global empty reachable = {}
+    loop
+        def frontier = source
+        def reachable = frontier
+        while (true)
+            def next_frontier = frontier . edge
+            def frontier = next_frontier
+            monus frontier = reachable    # frontier = frontier - reachable
+            upsert reachable = frontier   # reachable = reachable ∪ frontier
+            break break_reachable = empty(frontier)
+        end while
+    end loop
+end algorithm
+```
+
+The PyRel's metamodel IR for the above algorithm is constructed with the utilities as
+follows.
+
+```
+base_model = Model("algorithm_builder", dry_run=True)
+
+# Input (context) data
+
+edge = base_model.Relationship("Edge from {source:int} to {target:int}")
+source = base_model.Relationship("Source node {node:int}")
+
+with algorithm(base_model):
+    setup(
+        define(edge(1,2), edge(2,3), edge(3,4), edge(4,1))),
+        define(source(1))
+    )
+
+    # "local" variables and relations
+    n = Integer.ref()
+    m = Integer.ref()
+    reachable = base_model.Relationship("Reachable node {node:int}")
+    frontier = base_model.Relationship("Frontier node {node:int}")
+    next_frontier = base_model.Relationship("Next frontier node {node:int}")
+
+    global_(empty(define(reachable(n))))
+    assign(define(frontier(n)).where(source(n)))
+    assign(define(reachable(n)).where(frontier(n)))
+    with while_():
+        assign(define(next_frontier(m)).where(frontier(n), edge(n, m)))
+        assign(define(frontier(m)).where(next_frontier(m)))
+        monus(define(frontier(n)).where(reachable(n)))
+        upsert(0)(define(reachable(n)).where(frontier(n)))
+        break_(where(not_(frontier(n))))
+
+# Prints the PyRel Metamodel (IR)
+print(get_metamodel())
+
+# Prints the LQP transaction
+print(get_lqp_str())
+```
+"""
+from relationalai.semantics import Model
+from relationalai.semantics.metamodel import factory, ir, types
+from relationalai.semantics.internal.internal import Fragment
+from relationalai.semantics.lqp.algorithms import (
+    mk_empty, mk_assign, mk_upsert, mk_global, mk_monus
+)
+from relationalai.semantics.lqp.constructors import mk_transaction
+from relationalai.semantics.lqp.compiler import Compiler
+from relationalai.semantics.lqp import ir as lqp, builtins
+from typing import cast, TypeGuard, Optional, Sequence
+from lqp import print as lqp_print
+import threading
+from contextlib import contextmanager
+
+
+# While the constructors are very light-weight they enforce
+# the following grammar for algorithms:
+#
+# <Algorithm> := with algorithm(base_model): <Script>
+# <Script> := <Instruction>*
+# <Instruction> := <BaseInstruction> | <Loop>
+# <BaseInstruction> := [global_(] empty(Fragment) [)]
+#                    | [global_(] assign(<Fragment>) [)]
+#                    | break(<Fragment>)
+#                    | upsert(<Int>)(<Fragment>)
+#                    | monus(<Fragment>)
+# <Loop> := with while_(): <Script>
+#
+# Note: global_ annotation can only be used on top-level empty and assign instructions at the
+# top-level of the algorithm script.
+
+_storage = threading.local()
+
+def get_builder() -> 'AlgorithmBuilder':
+    """ Retrieves the thread-local AlgorithmBuilder instance."""
+    global _storage
+    if not(hasattr(_storage, "algorithm_builder")):
+        _storage.algorithm_builder = AlgorithmBuilder()
+    return _storage.algorithm_builder
+
+def get_metamodel() -> ir.Model:
+    """ Retrieves the compiled metamodel IR for the previous algorithm. Can only be used
+    after an algorithm has been defined."""
+    return get_builder().get_metamodel()
+
+def get_lqp_str() -> str:
+    """ Retrieves the LQP string representation for the previous algorithm. Can only be used
+    after an algorithm has been defined."""
+    return get_builder().get_lqp_str()
+
+@contextmanager
+def algorithm(model:Model):
+    """ Context manager for defining an algorithm on the given base model."""
+    get_builder().begin_algorithm(model)
+    yield
+    get_builder().end_algorithm()
+
+@contextmanager
+def while_():
+    """ Context manager for defining a while loop within an algorithm."""
+    get_builder().begin_while_loop()
+    yield
+    get_builder().end_while_loop()
+
+def setup(*stmts:Fragment):
+    """ Defines the setup section of an algorithm: a collection of PyRel statement that
+    prepare input data for the algorithm."""
+    builder = get_builder()
+    assert len(builder.script_stacks) == 1, "setup can only be called at the top-level of an algorithm"
+    assert builder.setup_fragments is None, "setup can only be called once per algorithm"
+    builder.set_setup_fragments(stmts)
+
+def global_(pos:int):
+    """ Marks a top-level `empty` or `assign` instruction as defining a global relation."""
+    assert type(pos) is int, "global_ can only be applied to empty and assign"
+    builder = get_builder()
+    assert len(builder.script_stacks) == 1, "global_ can only be applied to top-level instructions"
+    assert len(builder.script_stacks[0].instructions) == pos + 1
+    task = cast(ir.Task, mk_global(builder.script_stacks[0].instructions[pos]))
+    builder.script_stacks[0].instructions[pos] = task
+    builder.add_global_relation(task)
+
+def empty(stmt) -> int:
+    """ Marks a PyRel statement as an assignment of empty relation. The statement must not
+    have a body (no where clause)."""
+    assert has_empty_body(stmt), "Empty instruction must have an empty body"
+    task = get_builder().compile_statement(stmt)
+    task = cast(ir.Task, mk_empty(task))
+    return get_builder().append_task(task)
+
+def assign(stmt) -> int:
+    """ Marks a PyRel statement as an assignment instruction."""
+    task = get_builder().compile_statement(stmt)
+    task = cast(ir.Task, mk_assign(task))
+    return get_builder().append_task(task)
+
+def upsert_with_arity(arity:int, stmt:Fragment):
+    task = get_builder().compile_statement(stmt)
+    task = cast(ir.Task, mk_upsert(task, arity))
+    get_builder().append_task(task)
+
+def upsert(arity:int):
+    """ Marks a PyRel statement as an upsert instruction with the given arity."""
+    assert type(arity) is int and arity >= 0, "arity must be a non-negative integer"
+    return lambda stmt: upsert_with_arity(arity, stmt)
+
+def monus(stmt: Fragment) -> int:
+    """ Marks a PyRel statement as a Boolean monus (set difference) instruction."""
+    task = get_builder().compile_statement(stmt)
+    task = cast(ir.Task, mk_monus(task, types.Bool, "or", 0))
+    return get_builder().append_task(task)
+
+def break_(stmt):
+    """ Marks a PyRel statement as a break instruction. The statement must be headless (no define clause)."""
+    assert has_no_head(stmt), "Break instruction must have a headless fragment"
+    task = get_builder().compile_statement(stmt)
+    assert isinstance(task, ir.Logical)
+    break_condition = [cond for cond in task.body if not isinstance(cond, ir.Update)]
+    break_node = factory.break_(factory.logical(break_condition))
+    get_builder().append_task(break_node)
+
+def has_empty_body(stmt) -> TypeGuard[Fragment]:
+    if not isinstance(stmt, Fragment):
+        return False
+    return len(stmt._where) == 0
+
+def has_no_head(frag):
+    return len(frag._define) == 0
+
+
+class ScriptBuilder:
+    """
+    Builder for Loopy scripts.
+    """
+    def __init__(self):
+        self.instructions:list[ir.Task] = []
+
+    def add_task(self, instr:ir.Task) -> int:
+        self.instructions.append(instr)
+        return len(self.instructions) - 1
+
+    def build_script(self, annos:list[ir.Annotation]) -> ir.Sequence:
+        return factory.sequence(
+            tasks=self.instructions,
+            annos=[builtins.script_annotation()] + annos
+        )
+
+
+class AlgorithmBuilder:
+    """
+    Builder for Loopy algorithms.
+    """
+    def __init__(self):
+        self.script_stacks:list[ScriptBuilder] = []
+        self.compiled_model:Optional[ir.Model] = None
+        self.global_relations:list[str] = []
+        self.base_model:Optional[Model] = None
+        self.setup_fragments:Optional[list[Fragment]] = None
+
+    def begin_algorithm(self, base_model:Model):
+        self.base_model = base_model
+        self.script_stacks = [ScriptBuilder()]
+        self.compiled_model = None
+        self.global_relations:list[str] = []
+        self.setup_fragments:Optional[list[Fragment]] = None
+
+    def add_global_relation(self, task:ir.Task):
+        assert isinstance(task, ir.Logical)
+        for t in task.body:
+            if isinstance(t, ir.Update):
+                if t.relation.name not in self.global_relations:
+                    self.global_relations.append(t.relation.name)
+
+    def set_setup_fragments(self, fragments:Sequence[Fragment]):
+        self.setup_fragments = list(fragments)
+
+    def compile_statement(self, stmt:Fragment) -> ir.Task:
+        assert self.base_model is not None
+        task = self.base_model._compiler.compile_task(stmt)
+        return task
+
+    def append_task(self, task:ir.Task) -> int:
+        assert len(self.script_stacks) > 0
+        return self.script_stacks[-1].add_task(task)
+
+    def begin_while_loop(self):
+        script_builder = ScriptBuilder()
+        self.script_stacks.append(script_builder)
+
+    def end_while_loop(self):
+        script_builder = self.script_stacks.pop()
+        while_script = script_builder.build_script([builtins.while_annotation()])
+        loop = factory.loop(while_script, annos=[builtins.while_annotation()])
+        self.append_task(loop)
+
+    def end_algorithm(self):
+        assert len(self.script_stacks) == 1
+        script_builder = self.script_stacks.pop()
+        algorithm_script = script_builder.build_script([builtins.algorithm_annotation()])
+        setup = self.compile_setup()
+        algorithm_logical = factory.logical(setup + [algorithm_script])
+        self.compiled_model = factory.compute_model(algorithm_logical)
+
+    def compile_setup(self) -> list[ir.Logical]:
+        if self.setup_fragments is None:
+            return []
+        assert self.setup_fragments is not None
+        assert self.base_model is not None
+        setup_tasks = []
+        for stmt in self.setup_fragments:
+            task = self.base_model._compiler.compile_task(stmt)
+            setup_tasks.append(task)
+        return setup_tasks
+
+    def get_metamodel(self) -> ir.Model:
+        """ Retrieves the compiled metamodel IR for the previous algorithm. """
+        metamodel = self.compiled_model
+        assert metamodel is not None, "No metamodel available. You must first define algorithm."
+        return metamodel
+
+    def get_lqp_str(self) -> str:
+        lqp = self.get_lqp()
+        options = lqp_print.ugly_config.copy()
+        options[str(lqp_print.PrettyOptions.PRINT_NAMES)] = True
+        options[str(lqp_print.PrettyOptions.PRINT_DEBUG)] = False
+        lqp_str = lqp_print.to_string(lqp, options)
+        return lqp_str
+
+    def get_lqp(self):
+        model = self.get_metamodel()
+
+        compiler = Compiler()
+        rewritten_model = compiler.rewrite(model)
+        write_epoch = compiler.do_compile(rewritten_model, {'fragment_id': b"f1"})[1]
+
+        define = cast(lqp.Define, write_epoch.writes[0].write_type)
+        debug_info = define.fragment.debug_info
+
+        read_epoch = self._build_read_epoch(debug_info)
+
+        transaction = mk_transaction([write_epoch, read_epoch])
+
+        return transaction
+
+    def _build_read_epoch(self, debug_info:lqp.DebugInfo) -> lqp.Epoch:
+        reads = []
+
+        relation_id:dict[str,lqp.RelationId] = dict()
+        for rel_id, rel_name in debug_info.id_to_orig_name.items():
+            if rel_name in self.global_relations:
+                relation_id[rel_name] = rel_id
+
+        global_relation_names = [rel for rel in self.global_relations if rel in relation_id]
+
+        for (i, rel_name) in enumerate(global_relation_names):
+            read = lqp.Read(
+                meta = None,
+                read_type = lqp.Output(
+                    meta=None,
+                    name=f"{rel_name}",
+                    relation_id=relation_id[rel_name],
+                )
+            )
+            reads.append(read)
+
+        read_epoch = lqp.Epoch(
+            meta = None,
+            writes = [],
+            reads = reads,
+        )
+
+        return read_epoch

relationalai/semantics/tests/test_snapshot_abstract.py

@@ -20,7 +20,7 @@ class AbstractSnapshotTest(ABC):
     provider:Provider = cast(SFProvider, Provider()) # type: ignore
 
     def run_snapshot_test(self, snapshot, script_path, db_schema=None, use_sql=False, use_lqp=True, use_rel=False,
-                          use_direct_access=False, e2e=False, use_csv=True, e2e_only=False):
+                          use_direct_access=False, e2e=False, use_csv=True, e2e_only=False, emit_constraints=False):
         # Resolve use_lqp
         use_lqp = use_lqp and (not use_rel) # use_rel overrides because use_lqp is default.
 
@@ -47,6 +47,7 @@ class AbstractSnapshotTest(ABC):
                 'model_suffix': "" if not e2e else f"_{unique_name}",
                 'use_sql': use_sql,
                 'reasoner.rule.use_lqp': use_lqp,
+                'reasoner.rule.emit_constraints': emit_constraints,
                 'keep_model': False,
                 # fix the current time to keep snapshots stable
                 'datetime_now': datetime.datetime.fromisoformat("2025-12-01T12:00:00+00:00"),