relationalai 0.13.0.dev0__py3-none-any.whl → 0.13.2__py3-none-any.whl

relationalai/semantics/lqp/rewrite/extract_keys.py

@@ -0,0 +1,512 @@
+from __future__ import annotations
+
+from relationalai.semantics.metamodel import ir
+from relationalai.semantics.metamodel.compiler import Pass
+from relationalai.semantics.metamodel.visitor import Rewriter
+from relationalai.semantics.metamodel.util import OrderedSet
+from relationalai.semantics.metamodel import helpers, factory as f, types, builtins
+from typing import Optional, Any, Iterable, Tuple
+from collections import defaultdict
+
+"""
+Given an Output with a group of keys (some of them potentially null),
+* extract the lookups that bind (transitively) all the keys
+* extract the lookups that bind (transitively) properties of the keys
+* generate all the valid combinations of keys being present or not
+  * first all keys are present,
+  * then we remove one key at a time,
+  * then we remove two keys at a time,and so on.
+  * the last combination is when all the *nullable* keys are missing.
+* for each combination:
+  * create a compound (hash) key
+  * create a Logical that:
+      * contains all the relevant lookups for keys and properties
+      * contains negated lookups for null keys
+      * outputs using the compound key
+
+E.g., we go from
+
+Logical
+    Foo(foo)
+    rel1(foo, x)
+    Logical ^[v1=None]
+        rel2(foo, v1)
+    Logical ^[v2=None, k2=None]
+        rel3(foo, k2)
+        rel4(k2, v2)
+    Logical ^[v3=None, k3=None]
+        rel5(foo, y)
+        rel6(y, k3)
+        rel7(k3, v3)
+    output[foo, k2, k3](v1, v2, v3)
+
+to
+
+Logical
+    Logical
+        Logical
+            Foo(foo)
+            rel1(foo, x)
+            Logical ^[k2=None, k3=None]
+                rel3(foo, k2)
+                rel5(foo, y)
+                rel6(y, k3)
+            Logical ^[v1=None]
+                rel2(foo, v1)
+            Logical ^[v2=None, k2=None]
+                rel4(k2, v2)
+            Logical ^[v3=None, k3=None]
+                rel7(k3, v3)
+            construct(Hash, "Foo", foo, "Concept2", k2, "Concept3", k3, compound_key)
+            output[compound_key](v1, v2, v3)
+        Logical
+            Foo(foo)
+            rel1(foo, x)
+            Logical ^[k2=None, k3=None]
+                rel5(foo, y)
+                rel6(y, k3)
+            Not
+                Logical
+                    rel3(foo, k2)
+            Logical ^[v1=None]
+                rel2(foo, v1)
+            Logical ^[v3=None, k3=None]
+                rel7(k3, v3)
+            construct(Hash, "Foo", foo, "Concept3", k3, compound_key)
+            output[compound_key](v1, None, v3)
+        Logical
+            Foo(foo)
+            rel1(foo, x)
+            Logical ^[k2=None, k3=None]
+                rel3(foo, k2)
+                rel5(foo, y)
+            Not
+                Logical
+                    rel6(y, k3)
+            Logical ^[v1=None]
+                rel2(foo, v1)
+            Logical ^[v2=None, k2=None]
+                rel4(k2, v2)
+            construct(Hash, "Foo", foo, "Concept2", k2, compound_key)
+            output[compound_key](v1, v2, None)
+        Logical
+            Foo(foo)
+            rel1(foo, x)
+            Logical ^[k2=None, k3=None]
+                rel5(foo, y)
+            Not
+                Logical
+                    rel3(foo, k2)
+                    rel6(y, k3)
+            Logical ^[v1=None]
+                rel2(foo, v1)
+            construct(Hash, "Foo", foo, compound_key)
+            output[compound_key](v1, None, None)
+"""
+class ExtractKeys(Pass):
+    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
+        return ExtractKeysRewriter().walk(model)
+
+"""
+* First, figure out all tasks that are common for all alternative logicals that will be
+  generated
+* Second, generate the appropriate negated tasks for when some keys are missing
+* Third, add tasks that are needed for the final output columns
+  (e.g., lookup name if we select name)
+* Lastly, create the approprite Output task, with some of the columns being None if they are
+  missing (None will be filtered out in a later step -- we just need the column number to be
+  the same here).
+"""
+class ExtractKeysRewriter(Rewriter):
+    def __init__(self):
+        super().__init__()
+        self.compound_keys: dict[Any, ir.Var] = {}
+
+    def _get_compound_key(self, orig_keys: Iterable[ir.Var]) -> ir.Var:
+        if orig_keys in self.compound_keys:
+            return self.compound_keys[orig_keys]
+        compound_key = f.var("compound_key", types.Hash)
+        self.compound_keys[orig_keys] = compound_key
+        return compound_key
+
+    def handle_logical(self, node: ir.Logical, parent: ir.Node, ctx:Optional[Any]=None) -> ir.Logical:
+        outputs = [x for x in node.body if isinstance(x, ir.Output) and x.keys]
+        # We are not in a logical with an output at this level.
+        if not outputs:
+            new_body = self.walk_list(node.body, node)
+            return node if new_body is node.body else f.logical(new_body, node.hoisted)
+        assert len(outputs) == 1, "multiple outputs with keys in a logical"
+        output = outputs[0]
+        assert output.keys
+        output_keys = output.keys
+
+        all_vars = OrderedSet.from_iterable(output_keys)
+        all_vars.update(helpers.output_vars(output.aliases))
+
+        info = self.preprocess_logical(node, output_keys)
+        top_level_tasks, var_to_default, nullable_keys, non_nullable_keys = info
+        # we only need to transform the logical if there are nullable keys
+        if not nullable_keys:
+            return node
+
+        # Flatten the contents of (direct) logicals that comprise the body of this output
+        flat_body = OrderedSet[ir.Task]()
+        for task in node.body:
+            if isinstance(task, ir.Logical):
+                if any(isinstance(t, (ir.Aggregate, ir.Rank)) for t in task.body):
+                    flat_body.add(task)
+                else:
+                    flat_body.update(task.body)
+            elif isinstance(task, ir.Output):
+                continue
+            else:
+                flat_body.add(task)
+
+        # Add to all_vars the vars of lookups that were flattened above (from some logical)
+        for task in flat_body:
+            if isinstance(task, ir.Lookup):
+                all_vars.update(helpers.vars(task.args))
+
+        partitions, deps = self.partition_tasks(flat_body, all_vars)
+
+        # Compute all valid key combinations (keys that are not null)
+        combinations = self.key_combinations(nullable_keys, deps, 0, non_nullable_keys.get_list())
+        # there is no need to transform if there is only a single combination
+        if len(combinations) == 1:
+            return node
+
+        ####################################################################################
+
+        outer_body: list[ir.Task] = []
+        annos = list(output.annotations)
+        annos.append(f.annotation(builtins.output_keys, tuple(output_keys)))
+        # Create a compound key that will be used in place of the original keys.
+        compound_key = self._get_compound_key(output_keys)
+
+        for key_combination in combinations:
+            missing_keys = OrderedSet.from_iterable(output_keys)
+            missing_keys.difference_update(key_combination)
+
+            # top-level tasks are always present in every clone
+            body = OrderedSet.from_iterable(top_level_tasks)
+            # add tasks that are specific to the keys present in the current combination
+            for key in key_combination:
+                assert key in partitions
+                body.update(partitions[key])
+
+            # vars used in a positive context
+            positive_vars = OrderedSet[ir.Var]()
+            for task in body:
+                if isinstance(task, ir.Lookup):
+                    positive_vars.update(helpers.vars(task.args))
+
+            # handle the construct node in each clone
+            values: list[ir.Value] = [compound_key.type]
+            for key in output_keys:
+                if isinstance(key.type, ir.UnionType):
+                    # the typer can derive union types when multiple distinct entities flow
+                    # into a relation's field, so use AnyEntity as the type marker
+                    values.append(ir.Literal(types.String, "AnyEntity"))
+                else:
+                    assert isinstance(key.type, ir.ScalarType)
+                    values.append(ir.Literal(types.String, key.type.name))
+                if key in key_combination:
+                    values.append(key)
+            body.add(ir.Construct(None, tuple(values), compound_key, OrderedSet().frozen()))
+
+            # find variables used only inside the negated context
+            negative_vars = OrderedSet[ir.Var]()
+            for key in missing_keys:
+                negative_partition = partitions[key]
+                for task in negative_partition:
+                    if isinstance(task, ir.Lookup):
+                        args = helpers.vars(task.args)
+                        for arg in args:
+                            if arg not in positive_vars:
+                                negative_vars.add(arg)
+
+            out_vars = helpers.output_vars(output.aliases)
+            # output variables that depend on a missing key. They have to be skipped
+            missing_out_vars = OrderedSet[ir.Var]()
+            # output variables that depend on a negated variable that's not a missing key
+            # they have to be handled in a more complicated manner
+            problematic_out_vars = OrderedSet[ir.Var]()
+            for out_var in out_vars:
+                out_deps = deps[out_var]
+                if out_var in missing_keys:
+                    missing_out_vars.add(out_var)
+                elif any(x in missing_keys for x in out_deps):
+                    missing_out_vars.add(out_var)
+                elif any(x in negative_vars for x in out_deps):
+                    problematic_out_vars.add(out_var)
+
+            if problematic_out_vars:
+                assert len(problematic_out_vars) == 1
+                exclude_vars = deps[problematic_out_vars[0]] & negative_vars
+                has_problematic_var = True
+            else:
+                exclude_vars = out_vars
+                has_problematic_var = False
+
+            self.negate_missing_keys(body, missing_keys, var_to_default, partitions, deps,
+                                     out_vars, exclude_vars, negative_vars, has_problematic_var)
+
+            new_output_aliases = []
+            for alias, out_value in output.aliases:
+                if not isinstance(out_value, ir.Var):
+                    new_out_value = out_value
+                else:
+                    new_out_value = None if out_value in missing_out_vars else out_value
+                new_output_aliases.append((alias, new_out_value))
+            body.add(f.output(new_output_aliases, [compound_key], annos=annos))
+
+            # Create the final logical for this combination
+            outer_body.append(f.logical(tuple(body), []))
+
+        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]):
+        top_level_tasks = OrderedSet()
+        non_nullable_vars = OrderedSet()
+        nullable_vars = OrderedSet()
+        var_to_default: dict[ir.Var, ir.Default] = {}
+
+        for task in node.body:
+            # Top-level lookups use only non-nullable vars
+            if isinstance(task, ir.Lookup):
+                vars = helpers.vars(task.args)
+                non_nullable_vars.update(vars)
+                top_level_tasks.add(task)
+            elif isinstance(task, ir.Logical):
+                if self.noop_logical(task):
+                    top_level_tasks.add(task)
+                    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):
+                        non_nullable_vars.add(h)
+                    elif isinstance(h, ir.Default):
+                        var_to_default[h.var] = h
+                        # Hoisted vars with a non-None default are not nullable
+                        if h.value is not None:
+                            non_nullable_vars.add(h.var)
+                        else:
+                            nullable_vars.add(h.var)
+
+            # Variables appearing in an aggregate's group-by are not nullable
+            elif isinstance(task, ir.Aggregate):
+                top_level_tasks.add(task)
+                for v in task.group:
+                    non_nullable_vars.add(v)
+
+            elif isinstance(task, ir.Construct):
+                top_level_tasks.add(task)
+                non_nullable_vars.add(task.id_var)
+            elif isinstance(task, ir.Data):
+                top_level_tasks.add(task)
+                non_nullable_vars.update(task.vars)
+            # TODO: should Union and Match be used for hoisted vars?
+            elif isinstance(task, (ir.Not, ir.Match, ir.Union)):
+                top_level_tasks.add(task)
+
+            else:
+                if not isinstance(task, ir.Output):
+                    raise ValueError(f"Unexpected task type: {type(task)}")
+
+        # Any variable appearing in both sets is non-nullable
+        nullable_vars = nullable_vars - non_nullable_vars
+
+        non_nullable_keys = OrderedSet.from_iterable(output_keys)
+        non_nullable_keys.difference_update(nullable_vars)
+        nullable_keys = OrderedSet.from_iterable(output_keys)
+        nullable_keys.difference_update(non_nullable_keys)
+
+        return top_level_tasks, var_to_default, nullable_keys, non_nullable_keys
+
+    # given a set of variables, compute the tasks that each variable is using and also
+    # other variables needed for this variable to bind correctly
+    def partition_tasks(self, tasks:Iterable[ir.Task], vars:Iterable[ir.Var]):
+        partitions:dict[ir.Var, OrderedSet[ir.Task]] = defaultdict(OrderedSet)
+        dependencies:dict[ir.Var, OrderedSet[ir.Var]] = defaultdict(OrderedSet)
+
+        def dfs_collect_deps(task, deps):
+            if isinstance(task, ir.Lookup):
+                args = helpers.vars(task.args)
+                for i, v in enumerate(args):
+                    # v depends on all previous vars
+                    for j in range(i):
+                        deps[v].add(args[j])
+                    # for ternary+ lookups, a var also depends on the next vars
+                    if i > 0 and len(args) >= 3:
+                        for j in range(i+1, len(args)):
+                            deps[v].add(args[j])
+            elif isinstance(task, ir.Construct):
+                vars = helpers.vars(task.values)
+                for val_var in vars:
+                    deps[task.id_var].add(val_var)
+            elif isinstance(task, ir.Logical):
+                for child in task.body:
+                    dfs_collect_deps(child, deps)
+            elif isinstance(task, (ir.Match, ir.Union)):
+                for child in task.tasks:
+                    dfs_collect_deps(child, deps)
+
+        for task in tasks:
+            dfs_collect_deps(task, dependencies)
+
+        def dfs_transitive_deps(var, visited):
+            for dep_var in dependencies[var]:
+                if dep_var not in visited:
+                    visited.add(dep_var)
+                    dfs_transitive_deps(dep_var, visited)
+
+        transitive_deps = defaultdict(OrderedSet)
+        for var in list(dependencies.keys()):
+            visited = OrderedSet()
+            dfs_transitive_deps(var, visited)
+            transitive_deps[var] = visited
+        dependencies = transitive_deps
+
+        for var in vars:
+            extended_vars = OrderedSet[ir.Var]()
+            extended_vars.add(var)
+
+            there_is_progress = True
+            while there_is_progress:
+                there_is_progress = False
+                for task in tasks:
+                    if task in partitions[var]:
+                        continue
+
+                    if isinstance(task, (ir.Logical, ir.Match, ir.Union)):
+                        hoisted = helpers.hoisted_vars(task.hoisted)
+                        if var in hoisted:
+                            partitions[var].add(task)
+                            there_is_progress = True
+                    elif isinstance(task, ir.Construct):
+                        if task.id_var == var:
+                            partitions[var].add(task)
+                            there_is_progress = True
+                    elif isinstance(task, ir.Lookup):
+                        args = helpers.vars(task.args)
+                        if len(args) == 1 and args[0] in extended_vars:
+                            partitions[var].add(task)
+                            there_is_progress = True
+                        # NOTE: heuristics to have dot_joins work
+                        elif len(args) >= 3 and args[-2] in extended_vars:
+                            partitions[var].add(task)
+                            extended_vars.add(args[-1])
+                            there_is_progress = True
+                        elif len(args) > 1 and args[-1] in extended_vars:
+                            partitions[var].add(task)
+                            for arg in args[:-1]:
+                                extended_vars.add(arg)
+                            there_is_progress = True
+                    elif isinstance(task, ir.Not):
+                        if isinstance(task.task, ir.Logical):
+                            hoisted = helpers.hoisted_vars(task.task.hoisted)
+                            if var in hoisted:
+                                partitions[var].add(task)
+                                there_is_progress = True
+                    else:
+                        assert False, f"invalid node kind {type(task)}"
+
+        return partitions, dependencies
+
+    # Generate all the valid combinations of non-nullable keys and nullable keys.
+    def key_combinations(self, nullable_keys: OrderedSet[ir.Var], key_deps, idx: int, non_null_keys: list[ir.Var]) -> OrderedSet[Tuple[ir.Var]]:
+        if idx < len(nullable_keys):
+            key = nullable_keys[idx]
+            set1 = self.key_combinations(nullable_keys, key_deps, idx + 1, non_null_keys + [key])
+            set2 = self.key_combinations(nullable_keys, key_deps, idx + 1, non_null_keys)
+            set1.update(set2)
+            return set1
+        else:
+            final_keys = []
+            for k in non_null_keys:
+                # If a key depends on other keys, all of them should be present in this combination.
+                # If some dependency is not present, ignore the current key.
+                deps = key_deps.get(k)
+                if deps and any(dk in nullable_keys and dk not in non_null_keys for dk in deps):
+                    continue
+                final_keys.append(k)
+            return OrderedSet.from_iterable([tuple(final_keys)])
+
+    def negate_missing_keys(self, body, missing_keys, var_to_default, partitions, deps,
+                            out_vars, exclude_vars, negative_vars, has_problematic_var:bool):
+        # for keys that are not present in the current combination
+        # we have to include their tasks negated
+        negated_tasks = OrderedSet[ir.Task]()
+        positive_tasks = defaultdict(OrderedSet)
+
+        for key in missing_keys:
+            negative_body = OrderedSet[ir.Task]()
+            negative_partition = partitions[key]
+            for task in negative_partition:
+                # task is already present positively at the top-level context
+                if task in body:
+                    continue
+                if not isinstance(task, ir.Lookup):
+                    negative_body.add(task)
+                    continue
+
+                args = helpers.vars(task.args)
+                if args[-1] in exclude_vars and args[-1] not in missing_keys:
+                    positive_tasks[args[-1]].add(task)
+                else:
+                    negative_body.add(task)
+
+            negated_tasks.update(negative_body)
+            if len(negative_body) > 1:
+                body.add(f.not_(f.logical(tuple(negative_body), [])))
+            elif len(negative_body) == 1:
+                body.add(f.not_(negative_body[0]))
+
+        if positive_tasks:
+            for _, d in positive_tasks.items():
+                body.update(d)
+
+        for out_var in out_vars:
+            out_deps = deps[out_var]
+            if has_problematic_var and any(x in missing_keys for x in out_deps):
+                continue
+            elif not has_problematic_var and any(x in missing_keys or x in negative_vars for x in out_deps):
+                continue
+
+            default = var_to_default.get(out_var)
+            partition = OrderedSet.from_iterable(partitions[out_var])
+            partition.difference_update(body)
+            property_body = OrderedSet[ir.Task]()
+            if not has_problematic_var:
+                for task in partition:
+                    if isinstance(task, ir.Lookup):
+                        vars = helpers.vars(task.args)
+                        if any(x in negative_vars for x in vars):
+                            continue
+                    property_body.add(task)
+            else:
+                property_body.update(partition)
+            if property_body:
+                body.add(f.logical(tuple(property_body), [default] if default else []))

relationalai/semantics/lqp/rewrite/function_annotations.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+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):
+
+        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. Collects functional
+    dependencies from unique Require nodes and uses this information to identify functional
+    relations.
+    """
+
+    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
+
+
+class CollectFDsVisitor(Visitor):
+    """
+    Visitor collects all unique constraints.
+    """
+
+    # Currently, only information about k-functional fd is collected.
+    def __init__(self):
+        super().__init__()
+        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(:checked [, k])` annotation.
+    """
+
+    def __init__(self, functional_relations: dict[Relation, int]):
+        super().__init__()
+        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: 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: 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 | [function_annotation])
+        return node.reconstruct(node.engine, node.relation, node.args, node.effect, node.annotations)