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

relationalai/semantics/lqp/rewrite/extract_keys.py

@@ -1,577 +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 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 negative_vars:
-                    missing_out_vars.add(out_var)
-                elif 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 the body contains an aggregate, and the Logical hoists only the aggregate
-        # output, then this node behaves as a lookup
-        if any(isinstance(t, ir.Aggregate) for t in node.body):
-            hoisted_vars = helpers.hoisted_vars(node.hoisted)
-            aggregate_outputs = []
-            for t in node.body:
-                if isinstance(t, ir.Aggregate):
-                    aggregate_outputs.extend(
-                        v
-                        for v in helpers.vars(t.args)
-                        if not helpers.is_aggregate_input(v, t)
-                    )
-
-            if hoisted_vars == aggregate_outputs:
-                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 [out_var])
-                )