relationalai 0.13.5__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

relationalai/semantics/lqp/model2lqp.py

@@ -1,1056 +0,0 @@
-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, 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,
-)
-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 Sequence, Tuple, cast, Union, Optional
-from warnings import warn
-import re
-import uuid
-
-# Main access point for translating metamodel to lqp. Converts the model IR to an LQP epoch.
-def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[Optional[tuple], lqp.Epoch]:
-    assert_valid_input(model)
-    decls: list[lqp.Declaration] = []
-    reads: list[lqp.Read] = []
-
-    # LQP only accepts logical tasks
-    # These are asserted at init time
-    root = cast(ir.Logical, model.root)
-    for subtask in root.body:
-        assert isinstance(subtask, ir.Logical)
-        new_decls = _translate_to_decls(ctx, subtask)
-        decls.extend(new_decls)
-
-    # Add pyrel error attributes to reads.
-    for err_id in _extract_pyrel_error_ids(ctx, model):
-        ctx.output_ids.append(err_id)
-
-    reads.extend(_get_output_reads(ctx.output_ids))
-
-    export_info = None
-    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)
-
-    debug_info = lqp.DebugInfo(id_to_orig_name=ctx.rel_id_to_orig_name, meta=None)
-    fragment_id = lqp.FragmentId(id=fragment_name, meta=None)
-    fragment = lqp.Fragment(id=fragment_id, declarations=decls, meta=None, debug_info=debug_info)
-    define_op = lqp.Define(fragment=fragment, meta=None)
-
-    txn = lqp.Epoch(
-        reads=reads,
-        writes=[lqp.Write(write_type=define_op, meta=None)],
-        meta=None
-    )
-
-    return (export_info, txn)
-
-def _effect_bindings(effect: Union[ir.Output, ir.Update]) -> list[ir.Value]:
-    if isinstance(effect, ir.Output):
-        # Unions may not return anything. The generated IR contains a None value when this
-        # happens. We ignore it here.
-        return [v for v in helpers.output_values(effect.aliases) if v]
-    else:
-        return list(effect.args)
-
-def _get_output_reads(output_ids: list[tuple[lqp.RelationId, str]]) -> list[lqp.Read]:
-    reads = []
-    for (rel_id, name) in output_ids:
-        assert isinstance(rel_id, lqp.RelationId)
-        output = lqp.Output(name=name, relation_id=rel_id, meta=None)
-        reads.append(lqp.Read(read_type=output, meta=None))
-    return reads
-
-def _get_export_reads(descriptors: list[ExportDescriptor]) -> tuple[str, list, list[lqp.Read]]:
-    reads = []
-    csv_columns = []
-    col_info = []
-    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("-", "_")
-
-    # Note that the engine will append the transaction id to the export path for access
-    # control reasons. So the actual final path will be
-    # `.../export/{export_filename}/data_{txn_id}`
-    export_path = f"snowflake://APP_STATE.RAI_INTERNAL_STAGE/export/{export_filename}/data"
-    export_csv_config = lqp.ExportCSVConfig(
-        path=export_path,
-        data_columns=csv_columns,
-        compression="gzip",
-        partition_size=200,
-        syntax_escapechar='"',  # To follow Snowflake's expected format
-        meta=None,
-    )
-    reads.append(lqp.Read(read_type=lqp.Export(config=export_csv_config, meta=None), meta=None))
-    return (export_filename, col_info, reads)
-
-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)
-
-def _translate_to_constraint_decls(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
-    constraint_decls: list[lqp.Declaration] = []
-    for task in rule.body:
-        if isinstance(task, ir.Logical):
-            constraint_decls.extend(_translate_to_constraint_decls(ctx, task))
-            continue
-        else:
-            assert isinstance(task, ir.Require)
-            decl = _translate_to_constraint_decl(ctx, task)
-            if decl is not None:
-                constraint_decls.append(decl)
-    return constraint_decls
-
-def _translate_to_constraint_decl(ctx: TranslationCtx, rule: ir.Require) -> Optional[lqp.Declaration]:
-    fd = normalized_fd(rule)
-    assert fd is not None
-
-    # check for unresolved types
-    if any(types.is_any(var.type) for var in fd.keys + fd.values):
-        warn(f"Ignoring FD with unresolved type: {fd}")
-        return None
-
-    lqp_typed_keys = [_translate_term(ctx, key) for key in fd.keys]
-    lqp_typed_values = [_translate_term(ctx, value) for value in fd.values]
-    lqp_typed_vars:list[Tuple[lqp.Var, lqp.Type]] = lqp_typed_keys + lqp_typed_values # type: ignore
-    lqp_guard_atoms = [_translate_to_atom(ctx, atom) for atom in fd.guard]
-    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)
-
-    return lqp.FunctionalDependency(
-        name=lqp_name,
-        guard=lqp_guard,
-        keys=lqp_keys,
-        values=lqp_values,
-        meta=None
-    )
-
-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)
-
-    if len(effects) == 0:
-        return []
-
-    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."
-
-        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
-
-def _translate_to_standard_decl(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
-    effects = collect_by_type((ir.Output, ir.Update), rule)
-
-    # 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 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
-
-            # Convert literal arguments to LQP values
-            args = []
-            for a in annotation.args:
-                assert isinstance(a, ir.Literal)
-                args.append(mk_value(a.value))
-            attributes.append(mk_attribute(annotation.relation.name, args))
-    return attributes
-
-# Translates an effect (export, output, or update) into the corresponding def. Note that
-# this method only generates the def and not the LQP read operation, which is added later by
-# `_get_output_reads` and `_get_export_reads`.
-def _translate_effect(ctx: TranslationCtx, effect: Union[ir.Output, ir.Update], body: lqp.Formula) -> lqp.Declaration:
-    bindings = _effect_bindings(effect)
-
-    if isinstance(effect, ir.Output):
-        projection, eqs, suffix = _translate_output_bindings(ctx, bindings)
-        meta_id = effect.id
-
-        if helpers.is_export(effect):
-            def_name = "export_relation" + suffix
-        else:
-            def_name = "output" + suffix
-
-        def_name = ctx.output_names.get_name_by_id(meta_id, def_name)
-        rel_id = get_output_id(ctx, def_name, meta_id)
-    else:
-        projection, eqs = _translate_bindings(ctx, bindings)
-        def_name = effect.relation.name
-        rel_id = get_relation_id(ctx, effect.relation, projection)
-
-    eqs.append(body)
-    new_body = mk_and(eqs)
-
-    # Context bookkeeping for exports and outputs
-    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))
-        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))
-
-    # First we collect annotations on the effect itself, e.g. from something like
-    # `select(...).annotate(...)`.
-    annotations = effect.annotations
-    if isinstance(effect, ir.Update):
-        # Then we translate annotations on the relation itself, e.g.
-        # ```
-        # Bar.foo = model.Relationship(...)
-        # Bar.foo.annotate(...)
-        # ```
-        annotations = annotations | effect.relation.annotations
-
-    return lqp.Def(
-        name = rel_id,
-        body = mk_abstraction(projection, new_body),
-        attrs = _translate_annotations(annotations),
-        meta = None,
-    )
-
-def _translate_output_bindings(ctx: TranslationCtx, bindings: list[ir.Value]) -> Tuple[list[Tuple[lqp.Var, lqp.Type]], list[lqp.Formula], str]:
-    symbol_literals = []
-    non_symbols = []
-    for binding in bindings:
-        if isinstance(binding, ir.Literal) and binding.type == types.Symbol:
-            symbol_literals.append(binding.value)
-        else:
-            non_symbols.append(binding)
-    projection, eqs = _translate_bindings(ctx, non_symbols)
-    if len(symbol_literals) > 0:
-        name_suffix = "_"
-        name_suffix += "_".join(symbol_literals)
-    else:
-        name_suffix = ""
-
-    return projection, eqs, name_suffix
-
-def _translate_rank(ctx: TranslationCtx, rank: ir.Rank, body: lqp.Formula) -> lqp.Formula:
-    # Ascending rank is constructed using rel_primitive_sort. If a limit is added to an
-    # ascending rank we can use rel_primitive_top for an efficient evaluation.
-    #
-    # Descending rank is constructed as an ascending sort, then the ascending rank is
-    # subtracted from a count of the elements (plus 1 so that we still start at 1).
-    # Adding a limit to a descending rank is done by adding a filter of rank <= limit.
-
-    # Limits are the sort plus a filter on rank <= limit.
-    if all(o for o in rank.arg_is_ascending):
-        ascending = True
-    elif all(not o for o in rank.arg_is_ascending):
-        ascending = False
-    else:
-        raise Exception("Mixed orderings in rank are not supported yet.")
-
-    # Filter out the group-by variables, since they are introduced outside the rank.
-    input_args, input_eqs = _translate_bindings(ctx, list(rank.args))
-    introduced_meta_projs = [arg for arg in rank.projection if arg not in rank.group and arg not in rank.args]
-    projected_args, projected_eqs = _translate_bindings(ctx, list(introduced_meta_projs))
-
-    # rank expects an Int128 result, but the primitive will return an Int64 result.
-    # we need to set up an intermediary variable to hold the Int64 result, and a cast
-    # to convert it to Int128.
-    result_var, _ = _translate_term(ctx, rank.result)
-    # The primitive will return an Int64 result, so we need a var to hold the intermediary.
-    result_64_var = gen_unique_var(ctx, "v_rank")
-    result_64_type = mk_type(lqp.TypeName.INT)
-
-    cast = lqp.Cast(input=result_64_var, result=result_var, meta=None)
-
-    body = mk_and([body] + input_eqs + projected_eqs)
-    abstr_args = input_args + projected_args
-
-    if ascending:
-        ranker = _translate_ascending_rank(ctx, rank.limit, result_64_var, body, abstr_args)
-    else:
-        ranker = _translate_descending_rank(ctx, rank.limit, result_64_var, body, abstr_args)
-
-    return mk_exists([(result_64_var, result_64_type)], mk_and([ranker, cast]))
-
-def _translate_descending_rank(ctx: TranslationCtx, limit: int, result: lqp.Var, body: lqp.Formula, abstr_args) -> lqp.Formula:
-    result_var = result
-    result_type = mk_type(lqp.TypeName.INT)
-
-    # Rename abstracted args in the body to new variable names
-    var_map = {var.name: gen_unique_var(ctx, 't_' + var.name) for (var, _) in abstr_args}
-    body = utils.rename_vars_formula(body, var_map)
-    new_abstr_args = [(var_map[var.name], typ) for (var, typ) in abstr_args]
-
-    # Construct a conjunction of the ranking, a counter for the body, a subtraction
-    # of the rank from the count and an addition of 1. Wrap this in an abstraction.
-    count_res = gen_unique_var(ctx, "count_res")
-
-    # Add one to the count to account for the rank starting at 1.
-    one, one_eq = constant_to_var(ctx, to_lqp_value(1, types.Int64), "one")
-    one_bigger = gen_unique_var(ctx, "one_bigger")
-    addition = mk_primitive("rel_primitive_add_monotype", [count_res, one, one_bigger])
-
-    # Subtract the rank from the count + 1
-    asc_rank = gen_unique_var(ctx, "asc_rank")
-    subtraction = mk_primitive("rel_primitive_subtract_monotype", [one_bigger, result_var, asc_rank])
-
-    # Construct the ranking
-    desc_ranking_terms = [asc_rank] + [v[0] for v in abstr_args]
-    ranking = lqp.FFI(
-        meta=None,
-        name="rel_primitive_sort",
-        args=[mk_abstraction(new_abstr_args, body)],
-        terms=desc_ranking_terms,
-    )
-
-    # Count the number of rows in the body
-    count_type = meta_type_to_lqp(types.Int64)
-    count_var, count_eq = constant_to_var(ctx, to_lqp_value(1, types.Int64), "counter")
-    desc_body = mk_and([body, count_eq])
-    aggr_abstr_args = new_abstr_args + [(count_var, count_type)]
-    count_aggr = lqp.Reduce(
-        op=lqp_operator(
-            ctx,
-            "count",
-            "count",
-            mk_type(lqp.TypeName.INT)
-        ),
-        body=mk_abstraction(aggr_abstr_args, desc_body),
-        terms=[count_res],
-        meta=None
-    )
-
-    # Bring it all together and do the maths.
-    ranking = mk_exists(
-        vars=[
-            (asc_rank, result_type),
-            (count_res, result_type),
-            (one, result_type),
-            (one_bigger, result_type)
-        ],
-        value=mk_and([ranking, count_aggr, one_eq, addition, subtraction])
-    )
-
-    # If there is a limit, we need to add a filter to the ranking.
-    # Wrap with a rank <= limit
-    if limit != 0:
-        limit_term, _ = _translate_term(ctx, ir.Literal(types.Int64, limit))
-        limiter = mk_primitive("rel_primitive_lt_eq_monotype", [result_var, limit_term])
-        ranking = mk_and([ranking, limiter])
-
-    return ranking
-
-def _translate_ascending_rank(ctx: TranslationCtx, limit: int, result_var: lqp.Var, body: lqp.Formula, abstr_args) -> lqp.Formula:
-    terms = [result_var] + [v[0] for v in abstr_args]
-
-    # Rename abstracted args in the body to new variable names
-    var_map = {var.name: gen_unique_var(ctx, 't_' + var.name) for (var, _) in abstr_args}
-    body = utils.rename_vars_formula(body, var_map)
-    new_abstr_args = [(var_map[var.name], typ) for (var, typ) in abstr_args]
-    sort_abstr = mk_abstraction(new_abstr_args, body)
-
-    if limit == 0:
-        return lqp.FFI(
-            meta=None,
-            name="rel_primitive_sort",
-            args=[sort_abstr],
-            terms=terms,
-        )
-    else:
-        limit_type = meta_type_to_lqp(types.Int64)
-        limit_var, limit_eq = constant_to_var(ctx, to_lqp_value(limit, types.Int64), "limit")
-        limit_abstr = mk_abstraction([(limit_var, limit_type)], limit_eq)
-        return lqp.FFI(
-            meta=None,
-            name="rel_primitive_top",
-            args=[sort_abstr, limit_abstr],
-            terms=terms,
-        )
-
-def _rename_shadowed_abstraction_vars(
-    ctx: TranslationCtx,
-    aggr: ir.Aggregate,
-    abstr_args: list[Tuple[lqp.Var, lqp.Type]],
-    body_conjs: list[lqp.Formula]
-) -> list[Tuple[lqp.Var, lqp.Type]]:
-    """
-    Rename abstraction variables that shadow group-by variables.
-
-    This can happen when the same variable appears in both aggr.group and as an input
-    to the aggregation, e.g., min(Person.age).per(Person.age). The group-by variables
-    are in the outer scope, while the abstraction parameters are in the inner scope,
-    so we need different names to avoid shadowing.
-    """
-    # Get the LQP names of group-by variables
-    group_var_names = set()
-    for group_var in aggr.group:
-        lqp_var = _translate_var(ctx, group_var)
-        group_var_names.add(lqp_var.name)
-
-    # Rename any abstraction parameters that conflict with group-by variables
-    renamed_abstr_args = []
-    for var, typ in abstr_args:
-        if var.name in group_var_names:
-            # This variable shadows a group-by variable, so rename it
-            fresh_var = gen_unique_var(ctx, var.name)
-            # Add an equality constraint: fresh_var == var
-            # var is a free variable referring to the outer scope group-by variable
-            body_conjs.append(mk_primitive("rel_primitive_eq", [fresh_var, var]))
-            renamed_abstr_args.append((fresh_var, typ))
-        else:
-            renamed_abstr_args.append((var, typ))
-
-    return renamed_abstr_args
-
-def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Formula) -> Union[lqp.Reduce, lqp.Formula]:
-    # TODO: handle this properly
-    aggr_name = aggr.aggregation.name
-    supported_aggrs = ("sum", "count", "avg", "min", "max", "rel_primitive_solverlib_ho_appl")
-    assert aggr_name in supported_aggrs, f"only support {supported_aggrs} for now, not {aggr.aggregation.name}"
-
-    meta_output_terms = []
-    meta_input_terms = []
-
-    for (field, arg) in zip(aggr.aggregation.fields, aggr.args):
-        if field.input:
-            meta_input_terms.append(arg)
-        else:
-            meta_output_terms.append(arg)
-
-    output_terms = [_translate_term(ctx, term) for term in meta_output_terms]
-    output_vars = [term[0] for term in output_terms]
-
-    body_conjs = [body]
-    input_args, input_eqs = _translate_bindings(ctx, meta_input_terms)
-
-    # TODO: Can this safely be applied to all aggregates?
-    if aggr_name in ("sum", "min", "max"):
-        assert len(output_terms) == 1, f"{aggr_name} expects a single output variable"
-        assert len(meta_input_terms) == 1, f"{aggr_name} expects a single input variable"
-        assert isinstance(meta_output_terms[0], ir.Var)
-        assert input_args[0][1] == output_terms[0][1], f"{aggr_name}({input_args[0][1].type_name}) had output type of {output_terms[0][1].type_name}"
-
-    # Filter out the group-by variables, since they are introduced outside the aggregation.
-    # Input terms are added later below.
-    introduced_meta_projs = [arg for arg in aggr.projection if arg not in aggr.group and arg not in meta_input_terms]
-    projected_args, projected_eqs = _translate_bindings(ctx, list(introduced_meta_projs))
-    body_conjs.extend(input_eqs)
-    body_conjs.extend(projected_eqs)
-    abstr_args: list[Tuple[lqp.Var, lqp.Type]] = projected_args + input_args
-
-    # Rename abstraction variables that shadow group-by variables
-    abstr_args = _rename_shadowed_abstraction_vars(ctx, aggr, abstr_args, body_conjs)
-
-    if aggr_name == "count":
-        assert len(output_terms) == 1, "Count and avg expect a single output variable"
-        assert isinstance(meta_output_terms[0], ir.Var)
-        # Count sums up "1" for each row. We use the expected output type for the type
-        # of the count variable.
-        typ = meta_type_to_lqp(meta_output_terms[0].type)
-        one_var, eq = constant_to_var(ctx, to_lqp_value(1, meta_output_terms[0].type), "one")
-        body_conjs.append(eq)
-        abstr_args.append((one_var, typ))
-
-    # Average needs to wrap the reduce in Exists(Conjunction(Reduce, div))
-    if aggr_name == "avg":
-        assert len(output_vars) == 1, "avg should only have one output variable"
-        output_var = output_vars[0]
-
-        # Count sums up "1" for each row. We make the reasonably safe assumption that there
-        # are less than 2^31 rows in the body.
-        count_type = meta_type_to_lqp(types.Int64)
-        one_var, eq = constant_to_var(ctx, to_lqp_value(1, types.Int64), "one")
-        body_conjs.append(eq)
-        abstr_args.append((one_var, count_type))
-
-        # The average will produce two output variables: sum and count.
-        sum_result = gen_unique_var(ctx, "sum")
-        count_result = gen_unique_var(ctx, "count")
-
-        # Second to last is the variable we're summing over.
-        (sum_var, sum_type) = abstr_args[-2]
-        body = mk_and(body_conjs)
-
-        result = lqp.Reduce(
-            op=lqp_avg_op(ctx, aggr.aggregation.name, sum_var.name, sum_type),
-            body=mk_abstraction(abstr_args, body),
-            terms=[sum_result, count_result],
-            meta=None,
-        )
-
-        if sum_type == count_type:
-            div = mk_primitive("rel_primitive_divide_monotype", [sum_result, count_result, output_var])
-            conjunction = mk_and([result, div])
-
-            # Finally, we need to wrap everything in an `exists` to project away the sum and
-            # count variables and only keep the result of the division.
-            result_terms = [(sum_result, sum_type), (count_result, count_type)]
-        else:
-            # If the sum type and count type don't match, we need to cast the count
-            count_casted = gen_unique_var(ctx, "count_casted")
-            count_cast = lqp.Cast(input=count_result, result=count_casted, meta=None)
-
-            div = mk_primitive("rel_primitive_divide_monotype", [sum_result, count_casted, output_var])
-            conjunction = mk_and([result, count_cast, div])
-
-            # Finally, we need to wrap everything in an `exists` to project away the sum and
-            # count variables and only keep the result of the division.
-            result_terms = [(sum_result, sum_type), (count_result, count_type), (count_casted, sum_type)]
-
-        return mk_exists(result_terms, conjunction)
-
-    # `input_args` hold the types of the input arguments, but they may have been modified
-    # if we're dealing with a count, so we use `abstr_args` to find the type.
-    (aggr_arg, aggr_arg_type) = abstr_args[-1]
-    body = mk_and(body_conjs)
-
-    # Group-bys do not need to be handled at all, since they are introduced outside already
-    reduce = lqp.Reduce(
-        op=lqp_operator(ctx, aggr.aggregation.name, aggr_arg.name, aggr_arg_type),
-        body=mk_abstraction(abstr_args, body),
-        terms=output_vars,
-        meta=None
-    )
-    return reduce
-
-def _translate_to_formula(ctx: TranslationCtx, task: ir.Task) -> lqp.Formula:
-    if isinstance(task, ir.Logical):
-        # 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)
-    elif isinstance(task, ir.Not):
-        return lqp.Not(arg=_translate_to_formula(ctx, task.task), meta=None)
-    elif isinstance(task, ir.Exists):
-        lqp_vars, conjuncts = _translate_bindings(ctx, list(task.vars))
-        conjuncts.append(_translate_to_formula(ctx, task.task))
-        return mk_exists(lqp_vars, mk_and(conjuncts))
-    elif isinstance(task, ir.Construct):
-        assert len(task.values) >= 1, "Construct should have at least one value"
-        terms = [_translate_term(ctx, arg) for arg in task.values]
-        result_term = _translate_term(ctx, task.id_var)
-        terms.append(result_term)
-        assert result_term[1].type_name == lqp.TypeName.UINT128, \
-            f"Attempting to store a {task.id_var} in a type `{result_term[1].type_name}`"
-
-        return mk_primitive("rel_primitive_hash_tuple_uint128", [v for v, _ in terms])
-    elif isinstance(task, ir.Union):
-        disjs = [_translate_to_formula(ctx, child) for child in task.tasks]
-        return mk_or(disjs)
-    elif isinstance(task, (ir.Output, ir.Update)):
-        # Nothing to do here, handled in _translate_to_decls
-        return mk_and([])
-    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)}")
-
-# Only used for translating terms on relatoms, which can be specialized values.
-def _translate_relterm(ctx: TranslationCtx, term: ir.Value) -> Tuple[lqp.RelTerm, lqp.Type]:
-    if isinstance(term, ir.Literal) and term.type == types.Symbol:
-        if isinstance(term.value, str):
-            value = mk_value(term.value)
-            return mk_specialized_value(value), meta_type_to_lqp(types.String)
-        elif isinstance(term.value, int):
-            value = mk_value(term.value)
-            return mk_specialized_value(value), meta_type_to_lqp(types.Int64)
-        else:
-            raise NotImplementedError(f"Cannot specialize literal of type {type(term.value)}")
-    return _translate_term(ctx, term)
-
-def _translate_term(ctx: TranslationCtx, term: ir.Value) -> Tuple[lqp.Term, lqp.Type]:
-    if isinstance(term, ir.ScalarType):
-        # TODO: ScalarType is not like other terms, should be handled separately.
-        return to_lqp_value(term.name, types.String), meta_type_to_lqp(types.String)
-    elif isinstance(term, ir.Var):
-        t = meta_type_to_lqp(term.type)
-        return _translate_var(ctx, term), t
-    else:
-        assert isinstance(term, ir.Literal), f"Cannot translate value {term!r} of type {type(term)} to LQP Term; neither Var nor Literal."
-        v = to_lqp_value(term.value, term.type)
-        return v, meta_type_to_lqp(term.type)
-
-def _translate_to_atom(ctx: TranslationCtx, task: ir.Lookup) -> lqp.Formula:
-    if task.relation == builtins.cast:
-        assert len(task.args) == 3, f"expected three terms for {task.relation.name}, got {len(task.args)}"
-
-        terms = []
-        for arg in task.args[1:]:
-            term, _ = _translate_relterm(ctx, arg)
-            terms.append(term)
-        return lqp.Cast(input=terms[0], result=terms[1], meta=None)
-
-    if task.relation == builtins.join:
-        return _translate_join(ctx, task)
-    elif task.relation == builtins.infomap:
-        return _translate_infomap(ctx, task)
-    elif task.relation == builtins.louvain:
-        return _translate_louvain(ctx, task)
-    elif task.relation == builtins.label_propagation:
-        return _translate_label_propagation(ctx, task)
-
-    terms = []
-    term_types = []
-    for arg in task.args:
-        # Handle varargs, which come wrapped in a tuple.
-        if isinstance(arg, tuple):
-            for vararg in arg:
-                term, ty = _translate_relterm(ctx, vararg)
-                terms.append(term)
-                term_types.append(ty)
-        else:
-            term, ty = _translate_relterm(ctx, arg)
-            terms.append(term)
-            term_types.append(ty)
-
-    if builtins.is_pragma(task.relation):
-        lqp_name = pragma_to_lqp_name(task.relation.name)
-        return mk_pragma(lqp_name, terms)
-
-    if builtins.is_builtin(task.relation):
-        return build_primitive(task.relation.name, terms, term_types)
-
-    if helpers.is_external(task.relation):
-        return lqp.RelAtom(name=task.relation.name, terms=terms, meta=None)
-
-    projection, _ = _translate_bindings(ctx, list(task.args))
-    rid = get_relation_id(ctx, task.relation, projection)
-    return lqp.Atom(name=rid, terms=terms, meta=None)
-
-
-def get_relation_id(ctx: TranslationCtx, relation: ir.Relation, projection: list[Tuple[lqp.Var, lqp.Type]] = []) -> lqp.RelationId:
-    types = "_".join([str(t.type_name) for (_, t) in projection])
-    if types:
-        types = "_" + types
-    if relation.id in ctx.def_names.id_to_name:
-        unique_name = ctx.def_names.id_to_name[relation.id]
-    else:
-        name = helpers.relation_name_prefix(relation) + relation.name
-        name = helpers.sanitize(name)
-        unique_name = ctx.def_names.get_name_by_id(relation.id, name)
-
-    return utils.gen_rel_id(ctx, unique_name, types)
-
-def get_output_id(ctx: TranslationCtx, orig_name: str, metamodel_id: int) -> lqp.RelationId:
-    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: Sequence[ir.Value]) -> Tuple[list[Tuple[lqp.Var, lqp.Type]], list[lqp.Formula]]:
-    lqp_vars = []
-    conjuncts = []
-    for binding in bindings:
-        lqp_var, typ, eq = binding_to_lqp_var(ctx, binding)
-        lqp_vars.append((lqp_var, typ))
-        if eq is not None:
-            conjuncts.append(eq)
-
-    return lqp_vars, conjuncts
-
-def binding_to_lqp_var(ctx: TranslationCtx, binding: ir.Value) -> Tuple[lqp.Var, lqp.Type, Union[None, lqp.Formula]]:
-    if isinstance(binding, ir.Var):
-        var, typ = _translate_term(ctx, binding)
-        assert isinstance(var, lqp.Var)
-        return var, typ, None
-    elif isinstance(binding, ir.Literal):
-        lqp_value = to_lqp_value(binding.value, binding.type)
-        var, formula = constant_to_var(ctx, lqp_value)
-        return var, meta_type_to_lqp(binding.type), formula
-    else:
-        raise Exception(f"Unsupported binding type: {type(binding)}")
-
-def to_lqp_value(value: ir.PyValue, value_type: ir.Type) -> lqp.Value:
-    typ = meta_type_to_lqp(value_type)
-
-    # Ensure int values match the requested integer type.
-    if typ.type_name == lqp.TypeName.INT and isinstance(value, int):
-        assert lqp_types.INT_MIN <= value <= lqp_types.INT_MAX, f"{value} out of range for a 64-bit INT value"
-        val = value
-    elif typ.type_name == lqp.TypeName.INT128 and isinstance(value, int):
-        assert lqp_types.INT128_MIN <= value <= lqp_types.INT128_MAX, f"{value} out of range for an INT128 value"
-        val = lqp.Int128Value(value=value, meta=None)
-    elif typ.type_name == lqp.TypeName.UINT128 and isinstance(value, int):
-        assert lqp_types.UINT128_MIN <= value <= lqp_types.UINT128_MAX, f"{value} out of range for an UINT128 value"
-        val = lqp.UInt128Value(value=value, meta=None)
-    elif typ.type_name == lqp.TypeName.FLOAT and isinstance(value, float):
-        val = value
-    elif typ.type_name == lqp.TypeName.STRING and isinstance(value, str):
-        val = value
-    elif typ.type_name == lqp.TypeName.DECIMAL and isinstance(value, (int, float, PyDecimal)):
-        precision = typ.parameters[0].value
-        scale = typ.parameters[1].value
-        assert isinstance(precision, int) and isinstance(scale, int)
-        val = lqp.DecimalValue(precision=precision, scale=scale, value=PyDecimal(value), meta=None)
-    elif typ.type_name == lqp.TypeName.DATE and isinstance(value, date):
-        val = lqp.DateValue(value=value, meta=None)
-    elif typ.type_name == lqp.TypeName.DATETIME and isinstance(value, datetime):
-        utc_value = value.astimezone(timezone.utc) if value.tzinfo is not None else value # Convert to UTC cf. Iceberg
-        val = lqp.DateTimeValue(value=utc_value, meta=None)
-    elif typ.type_name == lqp.TypeName.BOOLEAN and isinstance(value, bool):
-        val = lqp.BooleanValue(value=value, meta=None)
-    else:
-        raise Exception(f"Unsupported type for constant: {typ.type_name} with value {value} of type {type(value)}")
-
-    return mk_value(val)
-
-def constant_to_var(ctx: TranslationCtx, value: lqp.Value, name_hint: str = "cvar") -> Tuple[lqp.Var, lqp.Formula]:
-    var = gen_unique_var(ctx, name_hint)
-    eq = mk_primitive("rel_primitive_eq", [var, value])
-    return var, eq
-
-def _extract_pyrel_error_ids(ctx: TranslationCtx, model: ir.Model) -> list[Tuple[lqp.RelationId, str]]:
-    effects = collect_by_type(ir.Update, model)
-    pyrel_error_attrs = []
-    # We use a separate counter here to avoid modifying the current model.
-    i = 1
-    for effect in effects:
-        if "pyrel_error_attrs" not in effect.relation.name:
-            continue
-        projection, _ = _translate_bindings(ctx, _effect_bindings(effect))
-        rel_id = get_relation_id(ctx, effect.relation, projection)
-        name = f"{effect.relation.name}_{i}"
-        pyrel_error_attrs.append((rel_id, name))
-        i += 1
-
-    return pyrel_error_attrs
-
-# Translate a relation reference into an abstraction over its fields.
-def _translate_relation_ref(ctx: TranslationCtx, relation: ir.Relation) -> lqp.Abstraction:
-    projection = []
-    for field in relation.fields:
-        var = gen_unique_var(ctx, field.name)
-        typ = meta_type_to_lqp(field.type)
-        projection.append((var, typ))
-
-    rid = get_relation_id(ctx, relation, projection)
-    atom = lqp.Atom(name=rid, terms=[var for (var, _) in projection], meta=None)
-    return mk_abstraction(projection, atom)
-
-# Common translation logic for graph algorithms.
-# task.args[0] : normalized weight list (int64, int64, float)
-# task.args[1] : normalized node count (relation or constant int64)
-# task.args[2] : normalized edge count (relation or constant int64)
-# task.args[3:-3] : algorithm parameters (var or constant)
-# task.args[-3] : diagnostic info
-# task.args[-2] : node index
-# task.args[-1] : community ident
-def _translate_graph_common(name: str, ctx: TranslationCtx, task: ir.Lookup):
-    abstractions = []
-
-    assert isinstance(task.args[0], ir.Relation), \
-        f"Expected relation as first arg to {name}, got {task.args[0]}:{type(task.args[0])}"
-    abstractions.append(_translate_relation_ref(ctx, task.args[0]))
-
-    # Allow constant args for node and edge count
-    for arg in task.args[1:3]:
-        if isinstance(arg, ir.Relation):
-            abst = _translate_relation_ref(ctx, arg)
-            typ = abst.vars[0][1]
-            assert typ.type_name == lqp.TypeName.INT, \
-                f"Expected Int64 types for node and edge counts, got type {typ.type_name}"
-            abstractions.append(abst)
-        else:
-            var, typ, eq = binding_to_lqp_var(ctx, arg)
-            assert eq is not None, \
-                f"Expected equality formula for {name} arg {arg}:{type(arg)}"
-            abstractions.append(mk_abstraction([(var, typ)], mk_and([eq])))
-
-    for arg in task.args[3:-3]:
-        var, typ, eq = binding_to_lqp_var(ctx, arg)
-        if eq:
-            abstractions.append(mk_abstraction([(var, typ)], mk_and([eq])))
-        else:
-            print(f"Primitive graph algorithm arg without eq:\n var:{var}, typ:{typ}\n arg:{arg}:{type(arg)}")
-            abstractions.append(mk_abstraction([(var, typ)], mk_and([])))
-
-    terms = []
-    for arg in task.args[-3:]:
-        term, _ = _translate_relterm(ctx, arg)
-        terms.append(term)
-
-    return (abstractions, terms)
-
-def _translate_infomap(ctx: TranslationCtx, task: ir.Lookup) -> lqp.Formula:
-    abstractions, terms = _translate_graph_common("infomap", ctx, task)
-
-    return lqp.FFI(
-        meta=None,
-        name="rel_primitive_infomap",
-        args=abstractions,
-        terms=terms,
-    )
-
-def _translate_louvain(ctx: TranslationCtx, task: ir.Lookup) -> lqp.Formula:
-    abstractions, terms = _translate_graph_common("louvain", ctx, task)
-
-    return lqp.FFI(
-        meta=None,
-        name="rel_primitive_louvain",
-        args=abstractions,
-        terms=terms,
-    )
-
-def _translate_label_propagation(ctx: TranslationCtx, task: ir.Lookup) -> lqp.Formula:
-    abstractions, terms = _translate_graph_common("label_propagation", ctx, task)
-
-    return lqp.FFI(
-        meta=None,
-        name="rel_primitive_async_label_propagation",
-        args=abstractions,
-        terms=terms,
-    )
-
-# Hard-coded implementation of Rel's string_join
-def _translate_join(ctx: TranslationCtx, task: ir.Lookup) -> lqp.Formula:
-    assert len(task.args) == 3
-    (strs, separator, target) = task.args
-    assert isinstance(separator, ir.Node)
-    assert isinstance(target, ir.Var)
-    assert isinstance(strs, tuple)
-
-    # Enumerate the strings we're going to be reducing over
-    enumerated_conjunctions = []
-    index_var, index_type = gen_unique_var(ctx, "idx"), meta_type_to_lqp(types.Int64)
-    string_var, string_type = gen_unique_var(ctx, "s"), meta_type_to_lqp(types.String)
-    for i, s in enumerate(strs):
-        index = to_lqp_value(i, types.Int64)
-        string = _translate_term(ctx, s)[0]
-        eq_idx = mk_primitive("rel_primitive_eq", [index_var, index])
-        eq_string = mk_primitive("rel_primitive_eq", [string_var, string])
-        enumerated_conjunctions.append(mk_and([eq_idx, eq_string]))
-    body = mk_abstraction([(index_var, index_type), (string_var, string_type)], mk_or(enumerated_conjunctions))
-
-    # Make the function we're reducing
-    a_var, a_type = gen_unique_var(ctx, "a"), meta_type_to_lqp(types.String)
-    b_var, b_type = gen_unique_var(ctx, "b"), meta_type_to_lqp(types.String)
-    curr_var, curr_type = gen_unique_var(ctx, "curr"), meta_type_to_lqp(types.String)
-    res_var, res_type = gen_unique_var(ctx, "res"), meta_type_to_lqp(types.String)
-
-    sep, _ = _translate_term(ctx, separator)
-
-    op = mk_abstraction(
-        [(a_var, a_type), (b_var, b_type), (curr_var, curr_type)],
-        mk_exists([(res_var, res_type)],
-            mk_and([
-                mk_primitive("rel_primitive_concat", [res_var, b_var, curr_var]),
-                mk_primitive("rel_primitive_concat", [a_var, sep, res_var])
-            ])
-        )
-    )
-
-    output_term = _translate_term(ctx, target)[0]
-
-    return lqp.Reduce(meta=None, op=op, body=body, terms=[output_term])
-
-def _translate_var(ctx: TranslationCtx, term: ir.Var) -> lqp.Var:
-    name = ctx.var_names.get_name_by_id(term.id, term.name)
-    return lqp.Var(name=name, meta=None)