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

relationalai/semantics/metamodel/rewrite/__init__.py

@@ -1,8 +0,0 @@
-from .discharge_constraints import DischargeConstraints
-from .dnf_union_splitter import DNFUnionSplitter
-from .extract_nested_logicals import ExtractNestedLogicals
-from .flatten import Flatten
-from .format_outputs import FormatOutputs
-from .handle_aggregations_and_ranks import HandleAggregationsAndRanks
-
-__all__ = ["DischargeConstraints", "DNFUnionSplitter", "ExtractNestedLogicals", "Flatten", "FormatOutputs", "HandleAggregationsAndRanks"]

relationalai/semantics/metamodel/rewrite/discharge_constraints.py

@@ -1,39 +0,0 @@
-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 rewrite_list
-
-
-class DischargeConstraints(c.Pass):
-    """
-    Since we should not generate code for `unique`, `exclusive`, `anyof` builtins all Require Check nodes marked with
-    `discharge` annotation when at least one of the builtins is presented in a check.
-    All Require/Check ir nodes marked with `discharge` annotation will be removed from the IR model in Flatten pass.
-    """
-
-    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:
-        return DischargeConstraintsVisitor().walk(model)
-
-
-@dataclass
-class DischargeConstraintsVisitor(v.Rewriter):
-    """
-    Visitor marks all nodes which should be removed from IR model with `discharge` annotation.
-    """
-
-    def handle_require(self, node: ir.Require, parent: ir.Node):
-        checks = rewrite_list(ir.Check, lambda n: self.walk(n, node), node.checks)
-        # discharge require if all the checks are discharged
-        if all(builtins.discharged_annotation in check.annotations for check in checks):
-            return node.reconstruct(node.engine, node.domain, checks, node.annotations | [builtins.discharged_annotation])
-        return node.reconstruct(node.engine, node.domain, checks, node.annotations)
-
-    def handle_check(self, node: ir.Check, parent: ir.Node):
-        check = self.walk(node.check, node)
-        assert isinstance(check, ir.Logical)
-        discharged_names = [builtins.unique.name, builtins.exclusive.name, builtins.anyof.name]
-        for item in check.body:
-            if isinstance(item, ir.Lookup) and item.relation.name in discharged_names:
-                return node.reconstruct(check, node.error, node.annotations | [builtins.discharged_annotation])
-        return node.reconstruct(check, node.error, node.annotations)

relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py

@@ -1,220 +0,0 @@
-from __future__ import annotations
-
-from relationalai.semantics.metamodel import ir
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel.visitor import Visitor, Rewriter
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
-from relationalai.semantics.metamodel import helpers, factory as f
-from typing import Optional, Any
-
-"""
-Handle DNF decomposition of unions
-
-E.g., we go from
-
-Logical
-    Logical
-        Union
-            Logical
-                Foo(foo::Foo)
-                a(foo::Foo, a::Int128)
-                a::Int128 < 100::Int128
-                Union
-                    Logical
-                        b(foo::Foo, b::Int128)
-                        b::Int128 < 100::Int128
-                    Logical
-                        c(foo::Foo, c::Int128)
-                        c::Int128 > 0::Int128
-            Logical
-                Foo(foo::Foo)
-                id(foo::Foo, id::Int128)
-                id::Int128 < 100::Int128
-            Logical
-                Foo(foo::Foo)
-                x(foo::Foo, x::Int128)
-                x::Int128 < 100::Int128
-                Union
-                    Logical
-                        y(foo::Foo, y::Int128)
-                        y::Int128 < 100::Int128
-                    Logical
-                        z(foo::Foo, z::Int128)
-                        z::Int128 > 0::Int128
-        ...
-        -> output(...)
-
-to
-
-Logical
-    Logical
-        Logical
-            Foo(foo::Foo)
-            a(foo::Foo, a::Int128)
-            a::Int128 < 100::Int128
-            Logical
-                b(foo::Foo, b::Int128)
-                b::Int128 < 100::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            a(foo::Foo, a::Int128)
-            a::Int128 < 100::Int128
-            Logical
-                c(foo::Foo, c::Int128)
-                c::Int128 > 0::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            id(foo::Foo, id::Int128)
-            id::Int128 < 100::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            x(foo::Foo, x::Int128)
-            x::Int128 < 100::Int128
-            Logical
-                y(foo::Foo, y::Int128)
-                y::Int128 < 100::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            x(foo::Foo, x::Int128)
-            x::Int128 < 100::Int128
-            Logical
-                z(foo::Foo, z::Int128)
-                z::Int128 > 0::Int128
-        ...
-        -> output(...)
-"""
-class DNFUnionSplitter(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        visitor = DNFExtractor()
-        model.accept(visitor)
-        return DNFRewriter(visitor).walk(model) if visitor.replaced_by else model
-
-class DNFExtractor(Visitor):
-    def __init__(self):
-        # The logical that contains the output.
-        # The assumption for the IR at this point is that there is only one output.
-        self.output_logical: Optional[ir.Logical] = None
-        self.output_keys: OrderedSet[ir.Var] = ordered_set()
-        self.active_negations: list[ir.Not] = []
-        # Nodes that have to split into multiple similar nodes, depending on the changes
-        # of sub-nodes.
-        self.should_split: OrderedSet[ir.Node] = OrderedSet()
-        # Track that a node has to be replaced by other nodes.
-        # If node X should be replaced by nodes Y and Z, then the parent should be replicated
-        # one time replacing X with Y and one time replacing X with Z.
-        self.replaced_by: dict[ir.Node, list[ir.Task]] = {}
-
-    def enter(self, node: ir.Node, parent: Optional[ir.Node]=None) -> Visitor:
-        if isinstance(node, ir.Logical):
-            if any(isinstance(x, ir.Output) for x in node.body):
-                assert not self.output_logical, "multiple outputs"
-                self.output_logical = node
-                output_node = next(x for x in node.body if isinstance(x, ir.Output))
-                self.output_keys = helpers.collect_vars(output_node)
-
-        elif isinstance(node, ir.Not):
-            self.active_negations.append(node)
-
-        return self
-
-    def leave(self, node: ir.Node, parent: Optional[ir.Node]=None) -> ir.Node:
-        if isinstance(node, ir.Logical) and node in self.should_split:
-            # The given logical, will be replaced by multiple logicals, each with a different
-            # group of tasks. We need to generate all the possible groups.
-            # A list of logical bodies (lists).
-            replacement_bodies: list[list[ir.Task]] = [[]]
-            for task in node.body:
-                if task in self.replaced_by:
-                    new_replacement_bodies: list[list[ir.Task]] = []
-                    replacement_tasks = self.replaced_by[task]
-                    for body in replacement_bodies:
-                        for new_task in replacement_tasks:
-                            # copy to mutate
-                            new_body = list(body)
-                            new_body.append(new_task.clone())
-                            new_replacement_bodies.append(new_body)
-                    replacement_bodies = new_replacement_bodies
-
-                else:
-                    for new_body in replacement_bodies:
-                        new_body.append(task.clone())
-
-            replacement_tasks: list[ir.Task] = []
-            for body in replacement_bodies:
-                new_task = f.logical(body, node.hoisted)
-                replacement_tasks.append(new_task)
-            self.replaced_by[node] = replacement_tasks
-
-            if node != self.output_logical:
-                self.should_split.add(parent)
-            elif node == self.output_logical:
-                assert parent and isinstance(parent, ir.Logical)
-                new_parent = f.logical(tuple(replacement_tasks), node.hoisted)
-                self.replaced_by[parent] = [new_parent]
-
-        elif isinstance(node, ir.Not) and self.active_negations[-1] == node:
-            self.active_negations.pop()
-
-        elif (isinstance(node, ir.Union) and
-              self.output_logical and
-              len(self.active_negations) % 2 == 0 and
-              len(node.tasks) > 1):
-            # We split the union when there are vars not present in all branches that are
-            # present in the output keys. If vars are not in output keys then they act as
-            # filters only and do not require splitting.
-            should_split = False
-            all_vars = helpers.collect_vars(node).get_set()
-            all_vars &= self.output_keys.get_set()
-
-            if all_vars:
-                for t in node.tasks:
-                    vars = helpers.collect_vars(t).get_set()
-                    curr_intersection = vars.intersection(all_vars)
-                    if curr_intersection != all_vars:
-                        should_split = True
-                        break
-
-            if should_split:
-                replacements:list[ir.Task] = []
-                for t in node.tasks:
-                    # If some branch should already be replaced, we flatten all
-                    # the replacements here.
-                    if t in self.replaced_by:
-                        replacements.extend(self.replaced_by[t])
-                    else:
-                        replacements.append(t)
-                self.replaced_by[node] = replacements
-                self.should_split.add(parent)
-
-        if isinstance(node, ir.Logical) and node == self.output_logical:
-            self.output_logical = None
-            self.output_keys = ordered_set()
-
-        return node
-
-class DNFRewriter(Rewriter):
-    def __init__(self, visitor: DNFExtractor):
-        super().__init__()
-        self.visitor = visitor
-        self.outer_parent_logical: Optional[ir.Logical] = None
-
-    def handle_logical(self, node: ir.Logical, parent: ir.Node, ctx:Optional[Any]=None) -> ir.Logical:
-        if node in self.visitor.replaced_by:
-            new_tasks = self.visitor.replaced_by[node]
-            assert len(new_tasks) == 1
-            new_task = new_tasks[0]
-            assert isinstance(new_task, ir.Logical)
-            return new_task
-        return node

relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py

@@ -1,78 +0,0 @@
-from __future__ import annotations
-
-from relationalai.semantics.metamodel import ir, factory as f, helpers
-from relationalai.semantics.metamodel.visitor import Rewriter, collect_by_type
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set, NameCache
-from relationalai.semantics.metamodel import dependency
-
-class ExtractNestedLogicals(Pass):
-
-    def __init__(self):
-        super().__init__()
-
-    #--------------------------------------------------
-    # Public API
-    #--------------------------------------------------
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        # process the root node
-        extractor = LogicalExtractor(model)
-        root = extractor.walk(model.root, model)
-
-        # no extractions, just return the original model
-        if not extractor.ctx.top_level:
-            return model
-
-        # create a new model with the extracted top level + the new root
-        body = extractor.ctx.top_level + [root]
-        return ir.Model(
-                model.engines,
-                OrderedSet.from_iterable(model.relations).update(extractor.ctx.relations).frozen(),
-                model.types,
-                ir.Logical(model.root.engine, tuple(), tuple(body))
-            )
-
-class LogicalExtractor(Rewriter):
-    def __init__(self, model):
-        super().__init__()
-        self.ctx = helpers.RewriteContext()
-        self.info = dependency.analyze(model.root)
-        self.name_cache = NameCache()
-
-    def handle_logical(self, node: ir.Logical, parent: ir.Node):
-        # rewrite the children
-        logical = super().handle_logical(node, parent)
-
-        # logicals that hoist vars and all vars do not have a default value will be
-        # extracted, except for logicals that require special treatment of the exposed
-        # variables (which is currently done by flatten), such as when the parent is a Match
-        # or a Union, of if the logical has a Rank.
-        if not (
-            logical.hoisted and
-            not isinstance(parent, (ir.Match, ir.Union)) and
-            all(isinstance(v, ir.Var) for v in logical.hoisted) and
-            not any(isinstance(c, ir.Rank) for c in logical.body)
-            ):
-            return logical
-
-        # compute the vars to be exposed by the extracted logical; those are keys (what
-        # makes the values unique) + the values (the hoisted variables)
-        exposed_vars = ordered_set()
-
-        # if there are aggregations, make sure we don't expose the projected and input vars,
-        # but expose groupbys
-        for agg in collect_by_type(ir.Aggregate, logical):
-            exposed_vars.difference_update(agg.projection)
-            exposed_vars.difference_update(helpers.aggregate_inputs(agg))
-            exposed_vars.update(agg.group)
-
-        # add the values (hoisted)
-        exposed_vars.update(helpers.hoisted_vars(logical.hoisted))
-
-        body = ordered_set()
-        body.update(self.info.task_dependencies(node)) # notice info is based on the original node
-        body.update(logical.body)
-
-        name = helpers.create_task_name(self.name_cache, logical, "_nested_logical")
-        connection = helpers.extract(logical, body, exposed_vars.get_list(), self.ctx, name)
-        return f.lookup(connection, exposed_vars.get_list())