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

relationalai/semantics/metamodel/rewrite/format_outputs.py

@@ -1,256 +0,0 @@
-from __future__ import annotations
-from typing import Tuple
-
-from relationalai.semantics.metamodel import builtins, ir, factory as f, types, visitor, helpers
-from relationalai.semantics.metamodel.compiler import Pass, group_tasks
-from relationalai.semantics.metamodel.util import OrderedSet
-from relationalai.semantics.metamodel.util import FrozenOrderedSet
-from relationalai.semantics.metamodel.typer.typer import is_primitive
-
-class FormatOutputs(Pass):
-    def __init__(self, use_rel: bool=False):
-        super().__init__()
-        self._use_rel = use_rel
-
-    #--------------------------------------------------
-    # Public API
-    #--------------------------------------------------
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        wide_outputs = options.get("wide_outputs", False)
-        return self.OutputRewriter(wide_outputs, self._use_rel).walk(model)
-
-    class OutputRewriter(visitor.Rewriter):
-        def __init__(self, wide_outputs: bool = False, use_rel: bool = False):
-            super().__init__()
-            self.wide_outputs = wide_outputs
-            self._use_rel = use_rel
-
-        def handle_logical(self, node: ir.Logical, parent: ir.Node):
-            # Rewrite children first
-            node = super().handle_logical(node, parent)
-
-            groups = group_tasks(node.body, {
-                "outputs": ir.Output,
-            })
-
-            # If no outputs, return as is
-            if not groups["outputs"]:
-                return node
-
-            if self.wide_outputs:
-                return adjust_wide_outputs(node, groups["outputs"])
-
-            return adjust_gnf_outputs(self._use_rel, node, groups["outputs"])
-
-#--------------------------------------------------
-# GNF vs wide output support
-#--------------------------------------------------
-
-# For wide outputs, only adjust the output task to include the keys.
-# output looks like: (key0, key1, val0, val1, ...)
-def adjust_wide_outputs(task: ir.Logical, outputs: OrderedSet[ir.Task]):
-    body = list(task.body)
-    for output in outputs:
-        assert(isinstance(output, ir.Output))
-        if output.keys:
-            body.remove(output)
-            body.append(rewrite_wide_output(output))
-    return ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations)
-
-# For GNF outputs we need to generate a rule for each "column" in the output
-# and potentially one wide key column
-def adjust_gnf_outputs(use_rel: bool, task: ir.Logical, outputs: OrderedSet[ir.Task]):
-    body = list(task.body)
-    for output in outputs:
-        assert(isinstance(output, ir.Output))
-        if output.keys:
-            # Remove the original output. This is replaced by per-column outputs below
-            body.remove(output)
-
-            is_export = helpers.is_export(output)
-
-            # Exports and Rel execution rely on all columns being in GNF format.
-            if is_export or use_rel:
-                _adjust_all_gnf_outputs(body, output, is_export)
-            else: # Otherwise, put all keys into one wide keys relation
-                _adjust_outputs_with_wide_keys(body, output)
-
-    return ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations)
-
-# Generate an output for each "column"
-# output looks like: def output(:cols, :col000, key0, key1, value)
-def _adjust_all_gnf_outputs(body, output: ir.Output, is_export: bool):
-    assert output.keys
-
-    original_cols = OrderedSet()
-    for idx, alias in enumerate(output.aliases):
-        # Skip None values which are used as a placeholder for missing values
-        if alias[1] is None:
-            continue
-        original_cols.add(alias[1])
-        body.extend(_generate_output_column_gnf(output, idx, alias, is_export))
-
-    idx = len(output.aliases)
-    for key in output.keys:
-        if key not in original_cols:
-            body.extend(_generate_output_column_gnf(output, idx, (key.name, key), is_export))
-            idx += 1
-
-# Generate an output for each value "column" and one wide output for all the keys
-#  * value output looks like: def output(:cols, :col000, key0, key1, value)
-#  * key output looks like:
-#    def output(:keys, output_key_0, output_key_1, other_key_0, ...)
-#
-# Exceptions: keys for exports and compound keys are converted to GNF, same as the
-# value columns.
-def _adjust_outputs_with_wide_keys(body, output: ir.Output):
-    assert output.keys
-
-    original_cols = OrderedSet()
-    val_cols: list[Tuple[str, ir.Value] | None] = []
-    key_cols: OrderedSet[Tuple[str, ir.Value]] = OrderedSet()
-    key_cols.add(("keys", f.literal("keys", types.Symbol))) # name key col so we can identify it later
-    for alias in output.aliases:
-        # None values are used as a placeholder for missing values
-        # They are added to maintain the correct col count when enumerated below
-        if alias[1] is None:
-            val_cols.append(None)
-            continue
-
-        original_cols.add(alias[1])
-
-        if isinstance(alias[1], ir.Var) and alias[1] in output.keys: # note: skips compound keys
-            key_cols.add(alias)
-        else:
-            val_cols.append(alias)
-
-    # Add keys not in output to the end
-    for key in output.keys:
-        if key not in original_cols:
-            key_cols.add((key.name, key))
-
-    # Generate GNF val cols
-    for idx, alias in enumerate(val_cols):
-        if alias:
-            new_col = _generate_output_column(output, idx, alias, key_cols)
-            body.extend(new_col)
-
-    # Create a wide key column with all keys
-    if len(key_cols) > 1:
-        body.append(ir.Output(
-            output.engine,
-            key_cols.frozen(),
-            output.keys,
-            output.annotations
-        ))
-
-# Generate a relation representing a single col in GNF form
-def _generate_output_column(output: ir.Output, idx: int, alias: tuple[str, ir.Value], key_cols):
-    if not output.keys:
-        return [output]
-
-    aliases = [("cols", f.literal("cols", types.Symbol))]
-    aliases.append(("col", f.literal(f"col{idx:03}", types.Symbol)))
-
-    # Append all keys at the start
-    keys = iter(key_cols)
-    assert next(keys) == ("keys", f.literal("keys", types.Symbol)) # skip col name
-    for key in keys:
-        aliases.append((f"key_{key[0]}_{idx}", key[1]))
-
-    aliases.append(alias) # append val
-
-    return [
-        ir.Output(
-            output.engine,
-            FrozenOrderedSet.from_iterable(aliases),
-            output.keys,
-            output.annotations
-        )
-    ]
-
-# Generate a relation representing a single col in GNF form for export
-def _generate_output_column_gnf(output: ir.Output, idx: int, alias: tuple[str, ir.Value], is_export: bool):
-    if not output.keys:
-        return [output]
-
-    aliases = [("cols", f.literal("cols", types.Symbol))] if not is_export else []
-    aliases.append(("col", f.literal(f"col{idx:03}", types.Symbol)))
-
-    for k in output.keys:
-        aliases.append((f"key_{k.name}_{idx}", k))
-
-    if (is_export and
-        isinstance(alias[1], ir.Var) and
-        (not is_primitive(alias[1].type) or alias[1].type == types.Hash)):
-
-        uuid = f.var(f"{alias[0]}_{idx}_uuid", types.String)
-
-        if not is_primitive(alias[1].type):
-            # For non-primitive types, we keep the original alias
-            aliases.append((alias[0], uuid))
-        else:
-            # For Hash types, we use the uuid name as alias
-            aliases.append((uuid.name, uuid))
-
-        return [
-            ir.Lookup(None, builtins.uuid_to_string, (alias[1], uuid)),
-            ir.Output(
-                output.engine,
-                FrozenOrderedSet.from_iterable(aliases),
-                output.keys,
-                output.annotations
-            )
-        ]
-    else:
-        aliases.append(alias)
-
-        return [
-            ir.Output(
-                output.engine,
-                FrozenOrderedSet.from_iterable(aliases),
-                output.keys,
-                output.annotations
-            )
-        ]
-
-def rewrite_wide_output(output: ir.Output):
-    assert(output.keys)
-
-    # Only append keys that are not already in the output
-    suffix_keys = []
-    for key in output.keys:
-        if all([val is not key for _, val in output.aliases]):
-            suffix_keys.append(key)
-
-    aliases: OrderedSet[Tuple[str, ir.Value]] = OrderedSet()
-
-    # Add the remaining args, unless it is already a key
-    for name, val in output.aliases:
-        if not isinstance(val, ir.Var) or val not in suffix_keys:
-            aliases.add((name, val))
-
-    # Add the keys to the output
-    for key in suffix_keys:
-        aliases.add((key.name, key))
-
-    # TODO - we are assuming that the Rel compiler will translate nullable lookups
-    # properly, returning a `Missing` if necessary, like this:
-    # (nested_192(_adult, _adult_name) or (not nested_192(_adult, _) and _adult_name = Missing)) and
-    return ir.Output(
-        output.engine,
-        aliases.frozen(),
-        output.keys,
-        output.annotations
-    )
-
-    # TODO: in the rel compiler, see if we can do this outer join
-    # 1. number of keys
-    # 2. each relation
-    # 3. each variable, starting with the keys
-    # 4. tag output with @arrow
-
-    # @arrow def output(_book, _book_title, _author_name):
-    #   rel_primitive_outer_join(#1, book_title, author_name, _book, _book_title, _author_name)
-    # def output(p, n, c):
-    #     rel_primitive_outer_join(#1, name, coolness, p, n, c)

relationalai/semantics/metamodel/rewrite/handle_aggregations_and_ranks.py

@@ -1,237 +0,0 @@
-from __future__ import annotations
-
-from relationalai.semantics.metamodel import ir, helpers
-from relationalai.semantics.metamodel.visitor import Rewriter
-from relationalai.semantics.metamodel.compiler import Pass, group_tasks
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
-from relationalai.semantics.metamodel import dependency
-
-# This rewrite pass handles aggregations and ranks to ensure that their dependencies are
-# self-contained in the required format for emitting Rel/LQP. The expected format is that
-# aggregations and ranks are each contained in their own Logical (which hoists the output
-# vars), with all dependencies pulled into the same Logical.
-#
-# For example,
-#
-# Logical ^[v::Int128]
-#     ... <dependencies> ...
-#     sum([foo::Foo], [], [a::Int128, v::Int128])
-#
-# Firstly, the pass ensures that all dependencies
-# required by an aggregation/rank are contained locally in the same Logical as the
-# aggregation/rank. For example, in the following Logical:
-#
-# Logical ^[v::Int128]
-#     Logical ^[a=None, foo=None]
-#         Foo(foo::Foo)
-#         a(foo::Foo, a::Int128)
-#     sum([foo::Foo], [], [a::Int128, v::Int128])
-# common(foo::Foo, a::Int128)
-#
-# The aggregation `sum` depends on the relation `common`. So, the lookup for `common` needs
-# to be pulled into the same Logical as the aggregation.
-#
-# Logical ^[v::Int128]
-#     Logical
-#         Logical ^[a=None, foo=None]
-#             Foo(foo::Foo)
-#             a(foo::Foo, a::Int128)
-#         common(foo::Foo, a::Int128)
-#     sum([foo::Foo], [], [a::Int128, v::Int128])
-# common(foo::Foo, a::Int128)
-#
-# Secondly, the pass separates Logicals containing more than one aggregation/rank into
-# separate Logicals, each containing a single aggregation/rank.
-#
-# Thirdly, the pass renames variables introduced inside aggregation/rank bodies to ensure
-# they do not clash with variables outside.
-
-class HandleAggregationsAndRanks(Pass):
-    def __init__(self):
-        super().__init__()
-
-    #--------------------------------------------------
-    # Public API
-    #--------------------------------------------------
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        dep_info = dependency.analyze(model.root)
-
-        r = AggregationsRanksRewriter(dep_info)
-        result = r.walk(model)
-
-        rn = AggregationsRanksVarRenameRewriter()
-        result = rn.walk(result)
-
-        return result
-
-# The AggregationsRanksRewriter ensures that each aggregation and rank is contained
-# in its own Logical, with all dependencies pulled into the same Logical.
-class AggregationsRanksRewriter(Rewriter):
-    def __init__(self, dep_info):
-        super().__init__()
-        self.info = dep_info
-        self.rewritten: dict[int, ir.Node] = {}
-
-    def handle_logical(self, node: ir.Logical, parent: ir.Node):
-        groups = group_tasks(node.body, {
-            "aggregates_and_ranks": (ir.Aggregate, ir.Rank),
-        })
-
-        aggregates_and_ranks = groups["aggregates_and_ranks"]
-
-        # If there are no aggregates or ranks, then just recurse into the Logical body
-        if not aggregates_and_ranks:
-            return super().handle_logical(node, parent)
-
-        agg_rank_logicals = []
-        for agg_rank in aggregates_and_ranks:
-            # Gather all dependencies of the logical containing the aggregate/rank
-            agg_deps = self.info.task_dependencies(agg_rank)
-
-            # Reconstruct the body of the Logical containing the aggregate/rank, starting
-            # with the existing body
-            body = ordered_set()
-
-            # agg_body is the inner body containing the dependencies of the aggregate/rank
-            agg_body = ordered_set()
-            for t in node.body:
-                if isinstance(t, (ir.Output, ir.Update)):
-                    # Outputs and Updates need to be kept in the outer body, rather than
-                    # nested inside the aggregate/rank body.
-                    body.add(t)
-                elif t not in aggregates_and_ranks:
-                    agg_body.add(t)
-
-            # Add all other dependencies
-            for dep in agg_deps:
-                # HACK: there are bugs in the dependency analysis that can cause cycles.
-                # Avoid these cycles because otherwise they can cause infinite recursion.
-                if agg_rank in self.info.task_dependencies(dep) or node in self.info.task_dependencies(dep):
-                    continue
-                agg_body.add(dep)
-
-            body.add(ir.Logical(node.engine, tuple(), tuple(agg_body)))
-
-            # Add the actual aggregate/rank
-            body.add(agg_rank.clone())
-
-            # Construct the final Logical holding the aggregate/rank contents.
-            # Output variables need to be hoisted
-            if isinstance(agg_rank, ir.Aggregate):
-                output_vars = [v for v in helpers.vars(agg_rank.args) if not helpers.is_aggregate_input(v, agg_rank)]
-            else:
-                assert isinstance(agg_rank, ir.Rank)
-                output_vars = [a for a in agg_rank.args]
-                output_vars.append(agg_rank.result)
-
-            agg_logical = ir.Logical(
-                engine=node.engine,
-                hoisted=tuple(output_vars),
-                body=tuple(body)
-            )
-            agg_rank_logicals.append(agg_logical)
-
-        if len(agg_rank_logicals) == 1:
-            # If there's only one, no need to create a parent Logical
-            result = agg_rank_logicals[0]
-        else:
-            # Otherwise, create a parent Logical, ensuring all body vars are hoisted
-            hoisted = OrderedSet()
-            for agg_rank_logical in agg_rank_logicals:
-                hoisted.update(agg_rank_logical.hoisted)
-
-            result = ir.Logical(
-                engine=node.engine,
-                hoisted=tuple(hoisted),
-                body=tuple(agg_rank_logicals)
-            )
-
-        # Rewrite the children
-        result = super().handle_logical(result, parent)
-
-        # Make a deep copy so that each task has a unique id. This is important for later
-        # rewrite passes (namely QuantifyVars) which identify tasks by id.
-        result = DeepCopyRewriter().walk(result)
-        return result
-
-# The AggregationsRanksVarRenameRewriter renames variables inside aggregation/rank
-# bodies to ensure they do not clash with variables outside. It is careful to keep certain
-# variables unrenamed because they need to interact with the outside, namely the group-by
-# variables and output variables.
-class AggregationsRanksVarRenameRewriter(Rewriter):
-    class RenameRewriter(Rewriter):
-        def __init__(self, to_keep: set[ir.Var], suffix: str):
-            super().__init__()
-            self.to_keep = to_keep
-            self.suffix = suffix
-            self.renamed_vars: dict[ir.Var, ir.Var] = {}
-
-        def handle_default(self, node: ir.Default, parent):
-            if node.var in self.to_keep:
-                return node
-
-            return ir.Default(
-                var=self.handle_var(node.var, node),
-                value=node.value
-            )
-
-        def handle_var(self, node: ir.Var, parent):
-            if node in self.to_keep:
-                return node
-
-            if node in self.renamed_vars:
-                return self.renamed_vars[node]
-
-            # Rename var
-            result = ir.Var(
-                type=node.type,
-                name=f"{node.name}_{self.suffix}"
-            )
-            self.renamed_vars[node] = result
-            return result
-
-    def __init__(self):
-        super().__init__()
-        self.renamed_vars: dict[ir.Var, ir.Var] = {}
-
-    def handle_logical(self, node: ir.Logical, parent: ir.Node):
-        groups = group_tasks(node.body, {
-            "aggregates_and_ranks": (ir.Aggregate, ir.Rank),
-        })
-
-        aggregates_and_ranks = groups["aggregates_and_ranks"]
-        if not aggregates_and_ranks:
-            return super().handle_logical(node, parent)
-
-        # There should only be one, because the AggregationsRanksRewriter should have
-        # separated them out.
-        assert len(aggregates_and_ranks) == 1, "Multiple aggregate/ranks still found after rewriting dependencies"
-
-        # Rename at this level
-        agg_rank = aggregates_and_ranks[0]
-        if isinstance(agg_rank, ir.Aggregate):
-            vars_to_keep = set(agg_rank.group)
-            output_vars = [v for v in helpers.vars(agg_rank.args) if not helpers.is_aggregate_input(v, agg_rank)]
-            vars_to_keep.update(output_vars)
-
-            result = self.RenameRewriter(vars_to_keep, 'agg').walk(node)
-        else:
-            assert isinstance(agg_rank, ir.Rank)
-            vars_to_keep = set(agg_rank.group)
-            output_vars = helpers.vars(agg_rank.args) + [agg_rank.result]
-            vars_to_keep.update(output_vars)
-            vars_to_keep.update(agg_rank.projection)
-
-            result = self.RenameRewriter(vars_to_keep, 'rank').walk(node)
-
-        # Process children
-        result = super().handle_logical(result, parent)
-        return result
-
-class DeepCopyRewriter(Rewriter):
-    def walk(self, node, parent=None):
-        new_node = super().walk(node, parent)
-        if isinstance(new_node, ir.Task):
-            return new_node.clone()
-        else:
-            return new_node