relationalai 0.12.8__py3-none-any.whl → 0.12.10__py3-none-any.whl

relationalai/semantics/lqp/executor.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from collections import defaultdict
+from datetime import datetime, timezone
 import atexit
 import re
 
@@ -13,6 +14,7 @@ from relationalai.semantics.lqp import result_helpers
 from relationalai.semantics.metamodel import ir, factory as f, executor as e
 from relationalai.semantics.lqp.compiler import Compiler
 from relationalai.semantics.lqp.intrinsics import mk_intrinsic_datetime_now
+from relationalai.semantics.lqp.constructors import mk_transaction
 from relationalai.semantics.lqp.types import lqp_type_to_sql
 from lqp import print as lqp_print, ir as lqp_ir
 from lqp.parser import construct_configure
@@ -39,6 +41,9 @@ class LQPExecutor(e.Executor):
         wide_outputs: bool = False,
         connection: Session | None = None,
         config: Config | None = None,
+        # In order to facilitate snapshot testing, we allow overriding intrinsic definitions
+        # like the current time, which would otherwise change between runs.
+        intrinsic_overrides: dict = {},
     ) -> None:
         super().__init__()
         self.database = database
@@ -48,6 +53,7 @@ class LQPExecutor(e.Executor):
         self.compiler = Compiler()
         self.connection = connection
         self.config = config or Config()
+        self.intrinsic_overrides = intrinsic_overrides
         self._resources = None
         self._last_model = None
         self._last_sources_version = (-1, None)
@@ -60,6 +66,8 @@ class LQPExecutor(e.Executor):
                 resource_class = rai.clients.snowflake.Resources
                 if self.config.get("use_direct_access", USE_DIRECT_ACCESS):
                     resource_class = rai.clients.snowflake.DirectAccessResources
+                if self.config.get("platform", "") == "local":
+                    resource_class = rai.clients.local.LocalResources
                 # NOTE: language="lqp" is not strictly required for LQP execution, but it
                 # will significantly improve performance.
                 self._resources = resource_class(
@@ -305,25 +313,28 @@ class LQPExecutor(e.Executor):
             config_dict['ivm.maintenance_level'] = lqp_ir.Value(value=ivm_flag, meta=None)
         return construct_configure(config_dict, None)
 
+    def _should_sync(self, model) :
+        if self._last_model != model:
+            return lqp_ir.Sync(fragments=[], meta=None)
+        else :
+            return None
+
     def _compile_intrinsics(self) -> lqp_ir.Epoch:
         """Construct an epoch that defines a number of built-in definitions used by the
         emitter."""
         with debugging.span("compile_intrinsics") as span:
             span["compile_type"] = "intrinsics"
 
+            now = self.intrinsic_overrides.get('datetime_now', datetime.now(timezone.utc))
+
             debug_info = lqp_ir.DebugInfo(id_to_orig_name={}, meta=None)
             intrinsics_fragment = lqp_ir.Fragment(
                 id = lqp_ir.FragmentId(id=b"__pyrel_lqp_intrinsics", meta=None),
-                declarations = [
-                    mk_intrinsic_datetime_now(),
-                ],
+                declarations = [mk_intrinsic_datetime_now(now)],
                 debug_info = debug_info,
                 meta = None,
             )
 
-
-            span["lqp"] = lqp_print.to_string(intrinsics_fragment, {"print_names": True, "print_debug": False, "print_csv_filename": False})
-
             return lqp_ir.Epoch(
                 writes=[
                     lqp_ir.Write(write_type=lqp_ir.Define(fragment=intrinsics_fragment, meta=None), meta=None)
@@ -331,6 +342,7 @@ class LQPExecutor(e.Executor):
                 meta=None,
             )
 
+    # [RAI-40997] We eagerly undefine query fragments so they are not committed to storage
     def _compile_undefine_query(self, query_epoch: lqp_ir.Epoch) -> lqp_ir.Epoch:
         fragment_ids = []
 
@@ -354,47 +366,40 @@ class LQPExecutor(e.Executor):
 
     def compile_lqp(self, model: ir.Model, task: ir.Task):
         configure = self._construct_configure()
+        # Merge the epochs into a single transaction. Long term the query bits should all
+        # go into a WhatIf action and the intrinsics could be fused with either of them. But
+        # for now we just use separate epochs.
+        epochs = []
+        epochs.append(self._compile_intrinsics())
 
-        model_txn = None
-        if self._last_model != model:
+        sync = self._should_sync(model)
+
+        if sync is not None:
             with debugging.span("compile", metamodel=model) as install_span:
                 install_span["compile_type"] = "model"
-                _, model_txn = self.compiler.compile(model, {"fragment_id": b"model"})
-                model_txn = txn_with_configure(model_txn, configure)
-                install_span["lqp"] = lqp_print.to_string(model_txn, {"print_names": True, "print_debug": False, "print_csv_filename": False})
+                _, model_epoch = self.compiler.compile(model, {"fragment_id": b"model"})
+                epochs.append(model_epoch)
                 self._last_model = model
 
-        with debugging.span("compile", metamodel=task) as compile_span:
-            compile_span["compile_type"] = "query"
+        with debugging.span("compile", metamodel=task) as txn_span:
             query = f.compute_model(f.logical([task]))
             options = {
                 "wide_outputs": self.wide_outputs,
                 "fragment_id": b"query",
             }
             result, final_model = self.compiler.compile_inner(query, options)
-            export_info, query_txn = result
-            query_txn = txn_with_configure(query_txn, configure)
-            compile_span["lqp"] = lqp_print.to_string(query_txn, {"print_names": True, "print_debug": False, "print_csv_filename": False})
+            export_info, query_epoch = result
 
-        # Merge the epochs into a single transactions. Long term the query bits should all
-        # go into a WhatIf action and the intrinsics could be fused with either of them. But
-        # for now we just use separate epochs.
-        epochs = []
+            epochs.append(query_epoch)
+            epochs.append(self._compile_undefine_query(query_epoch))
 
-        epochs.append(self._compile_intrinsics())
-
-        if model_txn is not None:
-            epochs.append(model_txn.epochs[0])
+            txn_span["compile_type"] = "query"
+            txn = mk_transaction(epochs=epochs, configure=configure, sync=sync)
+            txn_span["lqp"] = lqp_print.to_string(txn, {"print_names": True, "print_debug": False, "print_csv_filename": False})
 
-        query_txn_epoch = query_txn.epochs[0]
-        epochs.append(query_txn_epoch)
-        epochs.append(self._compile_undefine_query(query_txn_epoch))
-
-        txn = lqp_ir.Transaction(epochs=epochs, configure=configure, meta=None)
         validate_lqp(txn)
 
         txn_proto = convert_transaction(txn)
-        # TODO (azreika): Should export_info be encoded as part of the txn_proto? [RAI-40312]
         return final_model, export_info, txn_proto
 
     # TODO (azreika): This should probably be split up into exporting and other processing. There are quite a lot of arguments here...
@@ -462,12 +467,3 @@ class LQPExecutor(e.Executor):
             # If processing the results failed, revert to the previous model.
             self._last_model = previous_model
             raise e
-
-def txn_with_configure(txn: lqp_ir.Transaction, configure: lqp_ir.Configure) -> lqp_ir.Transaction:
-    """ Return a new transaction with the given configure. If the transaction already has
-    a configure, it is replaced. """
-    return lqp_ir.Transaction(
-        epochs=txn.epochs,
-        configure=configure,
-        meta=txn.meta,
-    )

relationalai/semantics/lqp/intrinsics.py

@@ -1,15 +1,16 @@
-from datetime import datetime, timezone
+from datetime import datetime
 
 from relationalai.semantics.lqp import ir as lqp
 from relationalai.semantics.lqp.constructors import mk_abstraction, mk_value, mk_type, mk_primitive
 from relationalai.semantics.lqp.utils import lqp_hash
 
-def mk_intrinsic_datetime_now() -> lqp.Def:
+# Constructs a definition of the current datetime.
+def mk_intrinsic_datetime_now(dt: datetime) -> lqp.Def:
     """Constructs a definition of the current datetime."""
     id = lqp_hash("__pyrel_lqp_intrinsic_datetime_now")
     out = lqp.Var(name="out", meta=None)
     out_type = mk_type(lqp.TypeName.DATETIME)
-    now = mk_value(lqp.DateTimeValue(value=datetime.now(timezone.utc), meta=None))
+    now = mk_value(lqp.DateTimeValue(value=dt, meta=None))
     datetime_now = mk_abstraction(
         [(out, out_type)],
         mk_primitive("rel_primitive_eq", [out, now]),

relationalai/semantics/lqp/ir.py

@@ -4,6 +4,7 @@ __all__ = [
     "SourceInfo",
     "LqpNode",
     "Declaration",
+    "FunctionalDependency",
     "Def",
     "Loop",
     "Abstraction",
@@ -45,6 +46,7 @@ __all__ = [
     "Read",
     "Epoch",
     "Transaction",
+    "Sync",
     "DebugInfo",
     "Configure",
     "IVMConfig",
@@ -59,6 +61,7 @@ from lqp.ir import (
     SourceInfo,
     LqpNode,
     Declaration,
+    FunctionalDependency,
     Def,
     Loop,
     Abstraction,
@@ -100,6 +103,7 @@ from lqp.ir import (
     Read,
     Epoch,
     Transaction,
+    Sync,
     DebugInfo,
     Configure,
     IVMConfig,

relationalai/semantics/lqp/model2lqp.py

@@ -11,7 +11,9 @@ from relationalai.semantics.lqp.constructors import (
 )
 from relationalai.semantics.lqp.utils import TranslationCtx, gen_unique_var
 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
@@ -19,8 +21,8 @@ from warnings import warn
 import re
 import uuid
 
-""" Main access point. Converts the model IR to an LQP transaction. """
-def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[Optional[tuple], lqp.Transaction]:
+# 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] = []
@@ -50,16 +52,10 @@ def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[
     fragment = lqp.Fragment(id=fragment_id, declarations=decls, meta=None, debug_info=debug_info)
     define_op = lqp.Define(fragment=fragment, meta=None)
 
-    txn = lqp.Transaction(
-        epochs=[
-            lqp.Epoch(
-                reads=reads,
-                writes=[lqp.Write(write_type=define_op, meta=None)],
-                meta=None
-            )
-        ],
-        configure=lqp.construct_configure({}, None),
-        meta=None,
+    txn = lqp.Epoch(
+        reads=reads,
+        writes=[lqp.Write(write_type=define_op, meta=None)],
+        meta=None
     )
 
     return (export_info, txn)
@@ -108,6 +104,43 @@ def _get_export_reads(export_ids: list[tuple[lqp.RelationId, int, lqp.Type]]) ->
     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)
+    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:
+        assert isinstance(task, ir.Require)
+        fd = normalized_fd(task)
+        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}")
+            continue
+
+        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
+
+        fd_decl = lqp.FunctionalDependency(
+            guard=lqp_guard,
+            keys=lqp_keys,
+            values=lqp_values,
+            meta=None
+        )
+
+        constraint_decls.append(fd_decl)
+
+    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)
@@ -458,7 +491,7 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
 
         return mk_exists(result_terms, conjunction)
 
-    # `input_args`` hold the types of the input arguments, but they may have been modified
+    # `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]
 

relationalai/semantics/lqp/passes.py

@@ -6,9 +6,11 @@ from relationalai.semantics.metamodel.util import FrozenOrderedSet
 
 from relationalai.semantics.metamodel.rewrite import Flatten
 
-from ..metamodel.rewrite import DischargeConstraints, DNFUnionSplitter, ExtractNestedLogicals, FormatOutputs
-from .rewrite import CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars, Splinter
-
+from ..metamodel.rewrite import DNFUnionSplitter, ExtractNestedLogicals, FormatOutputs
+from .rewrite import (
+    AnnotateConstraints, CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars,
+    Splinter, SplitMultiCheckRequires
+)
 from relationalai.semantics.lqp.utils import output_names
 
 from typing import cast, List, Sequence, Tuple, Union, Optional, Iterable
@@ -18,8 +20,9 @@ import hashlib
 
 def lqp_passes() -> list[Pass]:
     return [
+        SplitMultiCheckRequires(),
         FunctionAnnotations(),
-        DischargeConstraints(),
+        AnnotateConstraints(),
         Checker(),
         CDC(), # specialize to physical relations before extracting nested and typing
         ExtractNestedLogicals(), # before InferTypes to avoid extracting casts

relationalai/semantics/lqp/rewrite/__init__.py

@@ -1,15 +1,18 @@
+from .annotate_constraints import AnnotateConstraints
 from .cdc import CDC
 from .extract_common import ExtractCommon
 from .extract_keys import ExtractKeys
-from .function_annotations import FunctionAnnotations
+from .function_annotations import FunctionAnnotations, SplitMultiCheckRequires
 from .quantify_vars import QuantifyVars
 from .splinter import Splinter
 
 __all__ = [
+    "AnnotateConstraints",
     "CDC",
     "ExtractCommon",
     "ExtractKeys",
     "FunctionAnnotations",
     "QuantifyVars",
     "Splinter",
+    "SplitMultiCheckRequires",
 ]

relationalai/semantics/lqp/rewrite/annotate_constraints.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from relationalai.semantics.metamodel import builtins
+from relationalai.semantics.metamodel.ir import Node, Model, Require
+from relationalai.semantics.metamodel.compiler import Pass
+from relationalai.semantics.metamodel.rewrite.discharge_constraints import (
+    DischargeConstraintsVisitor
+)
+from relationalai.semantics.lqp.rewrite.functional_dependencies import (
+    is_valid_unique_constraint, normalized_fd
+)
+
+
+
+class AnnotateConstraints(Pass):
+    """
+    Extends `DischargeConstraints` pass by discharging only those Require nodes that cannot
+    be declared as constraints in LQP.
+
+    More precisely, the pass annotates Require nodes depending on how they should be
+    treated when generating code:
+     * `@declare_constraint` if the Require represents a constraint that can be declared in LQP.
+     * `@discharge` if the Require represents a constraint that should be dismissed during
+       code generation. Namely, when it cannot be declared in LQP and uses one of the
+       `unique`, `exclusive`, `anyof` builtins. These nodes are removed from the IR model
+       in the Flatten pass.
+    """
+
+    def rewrite(self, model: Model, options: dict = {}) -> Model:
+        return AnnotateConstraintsRewriter().walk(model)
+
+
+class AnnotateConstraintsRewriter(DischargeConstraintsVisitor):
+    """
+    Visitor marks all nodes which should be removed from IR model with `discharge` annotation.
+    """
+
+    def _should_be_declarable_constraint(self, node: Require) -> bool:
+        if not is_valid_unique_constraint(node):
+            return False
+        # Currently, we only declare non-structural functional dependencies.
+        fd = normalized_fd(node)
+        assert fd is not None  # already checked by _is_valid_unique_constraint
+        return not fd.is_structural
+
+    def handle_require(self, node: Require, parent: Node):
+        if self._should_be_declarable_constraint(node):
+            return node.reconstruct(
+                node.engine,
+                node.domain,
+                node.checks,
+                node.annotations | [builtins.declare_constraint_annotation]
+            )
+
+        return super().handle_require(node, parent)

relationalai/semantics/lqp/rewrite/extract_keys.py

@@ -249,6 +249,24 @@ class ExtractKeysRewriter(Rewriter):
 
         return f.logical(tuple(outer_body), [])
 
+    def noop_logical(self, node: ir.Logical) -> bool:
+        # logicals that don't hoist variables are essentially filters like lookups
+        if not node.hoisted:
+            return True
+        if len(node.body) != 1:
+            return False
+        inner = node.body[0]
+        if not isinstance(inner, (ir.Match, ir.Union)):
+            return False
+        outer_vars = helpers.hoisted_vars(node.hoisted)
+        inner_vars = helpers.hoisted_vars(inner.hoisted)
+        for v in outer_vars:
+            if v not in inner_vars:
+                return False
+        # all vars hoisted by the outer logical, are also
+        # hoisted by the inner Match/Union
+        return True
+
     # compute inital information that's needed for later steps. E.g., what's nullable or
     # not, do some output columns have a default value, etc.
     def preprocess_logical(self, node: ir.Logical, output_keys: Iterable[ir.Var]):
@@ -264,10 +282,11 @@ class ExtractKeysRewriter(Rewriter):
                 non_nullable_vars.update(vars)
                 top_level_tasks.add(task)
             elif isinstance(task, ir.Logical):
-                # logicals that don't hoist variables are essentially filters like lookups
-                if not task.hoisted:
+                if self.noop_logical(task):
                     top_level_tasks.add(task)
-                    # TODO: should we do something about the inner variables?
+                    non_nullable_vars.update(helpers.hoisted_vars(task.hoisted))
+                    continue
+
                 for h in task.hoisted:
                     # Hoisted vars without a default are not nullable
                     if isinstance(h, ir.Var):

relationalai/semantics/lqp/rewrite/function_annotations.py

@@ -1,79 +1,114 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from relationalai.semantics.metamodel import ir, compiler as c, visitor as v, builtins
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
+from typing import Optional
+from relationalai.semantics.metamodel import builtins
+from relationalai.semantics.metamodel.ir import (
+    Node, Model, Require, Logical, Relation, Annotation, Update
+)
+from relationalai.semantics.metamodel.compiler import Pass
+from relationalai.semantics.metamodel.visitor import Rewriter, Visitor
+from relationalai.semantics.lqp.rewrite.functional_dependencies import (
+    is_valid_unique_constraint, normalized_fd
+)
 
+# In the future iterations of PyRel metamodel, `Require` nodes will have a single `Check`
+# (and no `errors`). Currently, however, the unique constraints may result in multiple
+# `Check` nodes and for simplicity we split them in to separate `Require` nodes. This step
+# will be removed in the future.
+#
+# Note that unique constraints always have an empty `domain` so apply the splitting only
+# to such `Require` nodes.
+class SplitMultiCheckRequires(Pass):
+    """
+    Pass splits unique Require nodes that have empty domain but multiple checks into multiple
+    Require nodes with single check each.
+    """
+
+    def rewrite(self, model: Model, options: dict = {}) -> Model:
+        return SplitMultiCheckRequiresRewriter().walk(model)
+
+
+class SplitMultiCheckRequiresRewriter(Rewriter):
+    """
+    Splits unique Require nodes that have empty domain but multiple checks into multiple
+    Require nodes with single check each.
+    """
+    def handle_require(self, node: Require, parent: Node):
 
-class FunctionAnnotations(c.Pass):
+        if isinstance(node.domain, Logical) and not node.domain.body and len(node.checks) > 1:
+            require_nodes = []
+            for check in node.checks:
+                single_check = self.walk(check, node)
+                require_nodes.append(
+                    node.reconstruct(node.engine, node.domain, (single_check,), node.annotations)
+                )
+            return require_nodes
+
+        return node
+
+
+class FunctionAnnotations(Pass):
     """
-    Pass marks all appropriate relations with `function` annotation.
-    Criteria:
-    - there is a Require node with `unique` builtin (appeared as a result of `require(unique(...))`)
-    - `unique` declared for all the fields in a derived relation expect the last
+    Pass marks all appropriate relations with `function` annotation. Collects functional
+    dependencies from unique Require nodes and uses this information to identify functional
+    relations.
     """
 
-    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:
-        collect_fd = CollectFunctionalRelationsVisitor()
-        new_model = collect_fd.walk(model)
-        # mark relations collected by previous visitor with `@function` annotation
-        return FunctionalAnnotationsVisitor(collect_fd.functional_relations).walk(new_model)
+    def rewrite(self, model: Model, options: dict = {}) -> Model:
+        collect_fds = CollectFDsVisitor()
+        collect_fds.visit_model(model, None)
+        annotated_model = FunctionalAnnotationsRewriter(collect_fds.functional_relations).walk(model)
+        return annotated_model
 
 
-@dataclass
-class CollectFunctionalRelationsVisitor(v.Rewriter):
+class CollectFDsVisitor(Visitor):
     """
-    Visitor collects all relations which should be marked with `functional` annotation.
+    Visitor collects all unique constraints.
     """
 
+    # Currently, only information about k-functional fd is collected.
     def __init__(self):
         super().__init__()
-        self.functional_relations = ordered_set()
-
-    def handle_check(self, node: ir.Check, parent: ir.Node):
-        check = self.walk(node.check, node)
-        assert isinstance(check, ir.Logical)
-        unique_vars = []
-        for item in check.body:
-            # collect vars from `unique` builtin
-            if isinstance(item, ir.Lookup) and item.relation.name == builtins.unique.name:
-                var_set = set()
-                for vargs in item.args:
-                    assert isinstance(vargs, tuple)
-                    var_set.update(vargs)
-                unique_vars.append(var_set)
-        functional_rel = []
-        # mark relations as functional when at least 1 `unique` builtin
-        if len(unique_vars) > 0:
-            for item in check.body:
-                if isinstance(item, ir.Lookup) and not item.relation.name == builtins.unique.name:
-                    for var_set in unique_vars:
-                        # when unique declared for all the fields except the last one in the relation mark it as functional
-                        if var_set == set(item.args[:-1]):
-                            functional_rel.append(item.relation)
-
-        self.functional_relations.update(functional_rel)
-        return node.reconstruct(check, node.error, node.annotations)
-
-
-@dataclass
-class FunctionalAnnotationsVisitor(v.Rewriter):
+        self.functional_relations:dict[Relation, int] = {}
+
+    def visit_require(self, node: Require, parent: Optional[Node]):
+        if is_valid_unique_constraint(node):
+            fd = normalized_fd(node)
+            assert fd is not None
+            if fd.is_structural:
+                relation = fd.structural_relation
+                k = fd.structural_rank
+                current_k = self.functional_relations.get(relation, 0)
+                self.functional_relations[relation] = max(current_k, k)
+
+
+class FunctionalAnnotationsRewriter(Rewriter):
     """
-    This visitor marks functional_relations with `function` annotation.
+    This visitor marks functional_relations with `@function(:checked [, k])` annotation.
     """
 
-    def __init__(self, functional_relations: OrderedSet):
+    def __init__(self, functional_relations: dict[Relation, int]):
         super().__init__()
-        self._functional_relations = functional_relations
+        self.functional_relations = functional_relations
+
+    def get_functional_annotation(self, rel: Relation) -> Optional[Annotation]:
+        k = self.functional_relations.get(rel, None)
+        if k is None:
+            return None
+        if k == 1:
+            return builtins.function_checked_annotation
+        return builtins.function_ranked_checked_annotation(k)
 
-    def handle_relation(self, node: ir.Relation, parent: ir.Node):
-        if node in self._functional_relations:
-            return node.reconstruct(node.name, node.fields, node.requires, node.annotations | [builtins.function_checked_annotation],
-                                    node.overloads)
+    def handle_relation(self, node: Relation, parent: Node):
+        function_annotation = self.get_functional_annotation(node)
+        if function_annotation:
+            return node.reconstruct(node.name, node.fields, node.requires,
+                                    node.annotations | [function_annotation], node.overloads)
         return node.reconstruct(node.name, node.fields, node.requires, node.annotations, node.overloads)
 
-    def handle_update(self, node: ir.Update, parent: ir.Node):
-        if node.relation in self._functional_relations:
+    def handle_update(self, node: Update, parent: Node):
+        function_annotation = self.get_functional_annotation(node.relation)
+        if function_annotation:
             return node.reconstruct(node.engine, node.relation, node.args, node.effect,
-                                    node.annotations | [builtins.function_checked_annotation])
+                                    node.annotations | [function_annotation])
         return node.reconstruct(node.engine, node.relation, node.args, node.effect, node.annotations)