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

relationalai/semantics/lqp/rewrite/constants_to_vars.py

@@ -1,70 +0,0 @@
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel import ir, builtins as rel_builtins, factory as f, visitor
-from relationalai.semantics.metamodel.typer import typer
-
-from typing import List, Sequence, Tuple, Union
-
-# Rewrite constants to vars in Updates. This results in a more normalized format where
-# updates contain only variables. This allows for easier rewrites in later passes.
-class ConstantsToVars(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        r = self.ConstantToVarRewriter()
-        return r.walk(model)
-
-    # Return 1) a new list of Values with no duplicates (at the object level) and
-    # 2) equalities between any original Value and a deduplicated Value.
-    @staticmethod
-    def replace_constants_with_vars(vals: Sequence[ir.Value]) -> Tuple[List[ir.Value], List[ir.Lookup]]:
-        new_vals = []
-        eqs = []
-
-        for i, val in enumerate(vals):
-            if isinstance(val, ir.PyValue) or isinstance(val, ir.Literal):
-                # Replace constant with a new Var.
-                typ = typer.to_type(val)
-                assert isinstance(typ, ir.ScalarType), "can only replace scalar constants with vars"
-                new_var = ir.Var(typ, f"{typ.name.lower()}_{i}")
-                new_vals.append(new_var)
-                eqs.append(f.lookup(rel_builtins.eq, [new_var, val]))
-            else:
-                new_vals.append(val)
-
-        return new_vals, eqs
-
-    @staticmethod
-    def dedup_update(update: ir.Update) -> List[Union[ir.Update, ir.Lookup]]:
-        deduped_vals, req_lookups = ConstantsToVars.replace_constants_with_vars(update.args)
-        new_update = ir.Update(
-            update.engine,
-            update.relation,
-            tuple(deduped_vals),
-            update.effect,
-            update.annotations,
-        )
-        return req_lookups + [new_update]
-
-    # Does the actual work.
-    class ConstantToVarRewriter(visitor.Rewriter):
-        def __init__(self):
-            super().__init__()
-
-        # We implement handle_logical instead of handle_update because in
-        # addition to modifying said update we require new lookups (equality
-        # between original and deduplicated variables).
-        def handle_logical(self, node: ir.Logical, parent: ir.Node):
-            # In order to recurse over subtasks.
-            node = super().handle_logical(node, parent)
-
-            new_body = []
-            for subtask in node.body:
-                if isinstance(subtask, ir.Update):
-                    new_body.extend(ConstantsToVars.dedup_update(subtask))
-                else:
-                    new_body.append(subtask)
-
-            return ir.Logical(
-                node.engine,
-                node.hoisted,
-                tuple(new_body),
-                node.annotations
-            )

relationalai/semantics/lqp/rewrite/deduplicate_vars.py

@@ -1,104 +0,0 @@
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel import ir, builtins as rel_builtins, factory as f, visitor
-from relationalai.semantics.metamodel import helpers
-from relationalai.semantics.metamodel.util import FrozenOrderedSet
-
-from relationalai.semantics.lqp.utils import output_names
-
-from typing import List, Sequence, Tuple, Union
-
-# Deduplicate Vars in Updates and Outputs.
-class DeduplicateVars(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        r = self.VarDeduplicator()
-        return r.walk(model)
-
-    # Return 1) a new list of Values with no duplicates (at the object level) and
-    # 2) equalities between any original Value and a deduplicated Value.
-    @staticmethod
-    def dedup_values(vals: Sequence[ir.Value]) -> Tuple[List[ir.Value], List[ir.Lookup]]:
-        # If a var is seen more than once, it is a duplicate and we will create
-        # a new Var and equate it with the seen one.
-        seen_vars = set()
-
-        new_vals = []
-        eqs = []
-
-        for i, val in enumerate(vals):
-            # Duplicates can only occur within Vars.
-            # TODO: we don't know for sure if these are the only relevant cases.
-            if isinstance(val, ir.Default) or isinstance(val, ir.Var):
-                var = val if isinstance(val, ir.Var) else val.var
-                if var in seen_vars:
-                    new_var = ir.Var(var.type, var.name + "_dup_" + str(i))
-                    new_val = new_var if isinstance(val, ir.Var) else ir.Default(new_var, val.value)
-                    new_vals.append(new_val)
-                    eqs.append(f.lookup(rel_builtins.eq, [new_var, var]))
-                else:
-                    seen_vars.add(var)
-                    new_vals.append(val)
-            else:
-                # No possibility of problematic duplication.
-                new_vals.append(val)
-
-        return new_vals, eqs
-
-    # Returns a reconstructed output with no duplicate variable objects
-    # (dedup_values) and now necessary equalities between any two previously
-    # duplicate variables.
-    @staticmethod
-    def dedup_output(output: ir.Output) -> List[Union[ir.Output, ir.Lookup]]:
-        vals = helpers.output_values(output.aliases)
-        deduped_vals, req_lookups = DeduplicateVars.dedup_values(vals)
-        # Need the names so we can recombine.
-        alias_names = output_names(output.aliases)
-        new_output = ir.Output(
-            output.engine,
-            FrozenOrderedSet(list(zip(alias_names, deduped_vals))),
-            output.keys,
-            output.annotations,
-        )
-        return req_lookups + [new_output]
-
-    # Returns a replacement update with no duplicate variable objects
-    # (dedup_values) and now necessary equalities between any two previously
-    # duplicate variables.
-    @staticmethod
-    def dedup_update(update: ir.Update) -> List[Union[ir.Update, ir.Lookup]]:
-        deduped_vals, req_lookups = DeduplicateVars.dedup_values(update.args)
-        new_update = ir.Update(
-            update.engine,
-            update.relation,
-            tuple(deduped_vals),
-            update.effect,
-            update.annotations,
-        )
-        return req_lookups + [new_update]
-
-    # Does the actual work.
-    class VarDeduplicator(visitor.Rewriter):
-        def __init__(self):
-            super().__init__()
-
-        # We implement handle_logical instead of handle_update/handle_output
-        # because in addition to modifying said update/output we require new
-        # lookups (equality between original and deduplicated variables).
-        def handle_logical(self, node: ir.Logical, parent: ir.Node):
-            # In order to recurse over subtasks.
-            node = super().handle_logical(node, parent)
-
-            new_body = []
-            for subtask in node.body:
-                if isinstance(subtask, ir.Output):
-                    new_body.extend(DeduplicateVars.dedup_output(subtask))
-                elif isinstance(subtask, ir.Update):
-                    new_body.extend(DeduplicateVars.dedup_update(subtask))
-                else:
-                    new_body.append(subtask)
-
-            return ir.Logical(
-                node.engine,
-                node.hoisted,
-                tuple(new_body),
-                node.annotations
-            )

relationalai/semantics/lqp/rewrite/eliminate_data.py

@@ -1,108 +0,0 @@
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel import ir, builtins as rel_builtins, factory as f, visitor
-
-from typing import cast
-import pandas as pd
-import hashlib
-
-# Creates intermediary relations for all Data nodes and replaces said Data nodes
-# with a Lookup into these created relations. Reuse duplicate created relations.
-class EliminateData(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        r = self.DataRewriter()
-        return r.walk(model)
-
-    # Does the actual work.
-    class DataRewriter(visitor.Rewriter):
-        new_relations: list[ir.Relation]
-        new_updates: list[ir.Logical]
-        # Counter for naming new relations.
-        # It must be that new_count == len new_updates == len new_relations.
-        new_count: int
-        # Cache for Data nodes to avoid creating duplicate intermediary relations
-        data_cache: dict[str, ir.Relation]
-
-        def __init__(self):
-            self.new_relations = []
-            self.new_updates = []
-            self.new_count = 0
-            self.data_cache = {}
-            super().__init__()
-
-        # Create a cache key for a Data node based on its structure and content
-        def _data_cache_key(self, node: ir.Data) -> str:
-            values = pd.util.hash_pandas_object(node.data).values
-            return hashlib.sha256(bytes(values)).hexdigest()
-
-        def _intermediary_relation(self, node: ir.Data) -> ir.Relation:
-            cache_key = self._data_cache_key(node)
-            if cache_key in self.data_cache:
-                return self.data_cache[cache_key]
-            self.new_count += 1
-            intermediary_name = f"formerly_Data_{self.new_count}"
-
-            intermediary_relation = f.relation(
-                intermediary_name,
-                [f.field(v.name, v.type) for v in node.vars]
-            )
-            self.new_relations.append(intermediary_relation)
-
-            intermediary_update = f.logical([
-                # For each row (union), equate values and their variable (logical).
-                f.union(
-                    [
-                        f.logical(
-                            [
-                                f.lookup(rel_builtins.eq, [f.literal(val, var.type), var])
-                                for (val, var) in zip(row, node.vars)
-                            ],
-                        )
-                        for row in node
-                    ],
-                    hoisted = node.vars,
-                ),
-                # And pop it back into the relation.
-                f.update(intermediary_relation, node.vars, ir.Effect.derive),
-            ])
-            self.new_updates.append(intermediary_update)
-
-            # Cache the result for reuse
-            self.data_cache[cache_key] = intermediary_relation
-
-            return intermediary_relation
-
-        # Create a new intermediary relation representing the Data (and pop it in
-        # new_updates/new_relations) and replace this Data with a Lookup of said
-        # intermediary.
-        def handle_data(self, node: ir.Data, parent: ir.Node) -> ir.Lookup:
-            intermediary_relation = self._intermediary_relation(node)
-            replacement_lookup = f.lookup(intermediary_relation, node.vars)
-
-            return replacement_lookup
-
-        # Walks the model for the handle_data work then updates the model with
-        # the new state.
-        def handle_model(self, model: ir.Model, parent: None):
-            walked_model = super().handle_model(model, parent)
-            assert len(self.new_relations) == len(self.new_updates) and self.new_count == len(self.new_relations)
-
-            # This is okay because its LQP.
-            assert isinstance(walked_model.root, ir.Logical)
-            root_logical = cast(ir.Logical, walked_model.root)
-
-            # We may need to add the new intermediaries from handle_data to the model.
-            if self.new_count  == 0:
-                return model
-            else:
-                return ir.Model(
-                    walked_model.engines,
-                    walked_model.relations | self.new_relations,
-                    walked_model.types,
-                    ir.Logical(
-                        root_logical.engine,
-                        root_logical.hoisted,
-                        root_logical.body + tuple(self.new_updates),
-                        root_logical.annotations,
-                    ),
-                    walked_model.annotations,
-                )

relationalai/semantics/lqp/rewrite/extract_common.py

@@ -1,340 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Optional
-from relationalai.semantics.metamodel import ir, factory as f, helpers, visitor
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
-from relationalai.semantics.metamodel import dependency
-from relationalai.semantics.metamodel import builtins
-
-from functools import reduce
-
-class ExtractCommon(Pass):
-    """
-    Pass to analyze Logical bodies and extract lookups in their own Logical if it makes
-    sense. The heuristic is that, if there are multiple lookups, and there are also multiple
-    sibling nested logicals that will eventually be extracted by Flatten, then it makes
-    sense to extract these logicals into their own "rule", and then make the original body
-    just lookup this common rule.
-
-    From:
-        Logical
-            Logical
-                lookup1
-                lookup2
-                Logical1 ...
-                Logical2 ...
-    To:
-        Logical
-            Logical
-                lookup1
-                lookup2
-                derive common
-            Logical
-                lookup common
-                Logical1 ...
-                Logical2 ...
-    """
-
-    # The extraction plan heuristic is as follows:
-    #
-    # Given a set of binder tasks B and a set of extractable tasks E, we find:
-    #   - A subset of common tasks C in B, and
-    #   - A subset of exposed variables V output from tasks in C
-    # where:
-    #   - The intersection of common dependencies of all tasks in E are contained in C
-    #     (including transitive dependencies)
-    #   - The union of input variables for all tasks in E intersected with the output
-    #     variables of tasks in C are contained in V
-
-    #--------------------------------------------------
-    # Public API
-    #--------------------------------------------------
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        # create the  context
-        ctx = ExtractCommon.Context(model, options)
-
-        # rewrite the root
-        replacement = self.handle(model.root, ctx)
-
-        # the new root contains the extracted top level logicals and the rewritten root
-        if ctx.rewrite_ctx.top_level:
-            new_root = ir.Logical(model.root.engine, tuple(), tuple(ctx.rewrite_ctx.top_level + [replacement]))
-        else:
-            new_root = replacement
-
-        # create the new model, updating relations and root
-        return ir.Model(
-            model.engines,
-            OrderedSet.from_iterable(model.relations).update(ctx.rewrite_ctx.relations).frozen(),
-            model.types,
-            new_root
-        )
-
-    #--------------------------------------------------
-    # Extra classes
-    #--------------------------------------------------
-
-    class Context():
-        def __init__(self, model: ir.Model, options: dict):
-            self.rewrite_ctx = helpers.RewriteContext()
-            self.info = dependency.analyze(model.root)
-            self.options = options
-
-    @dataclass
-    class ExtractionPlan():
-        # tasks to extract to the body of the common logical
-        common_body: OrderedSet[ir.Task]
-        # tasks to remain in the original body
-        remaining_body: OrderedSet[ir.Task]
-        # variables to be exposed by the common logical
-        exposed_vars: OrderedSet[ir.Var]
-        # map from nested composite to the tasks in the common body that still need to be
-        # included in its body, because it contains variables not exposed by the common logical
-        local_dependencies: dict[ir.Task, OrderedSet[ir.Task]]
-        # a reference to the common connection created for this plan, if any
-        common_reference: Optional[ir.Lookup] = None
-
-    #--------------------------------------------------
-    # IR handlers
-    #--------------------------------------------------
-
-    def handle(self, task: ir.Task, ctx: Context):
-        # Currently we only extract if it's a sequence of Logicals, but we could in the
-        # future support other intermediate nodes
-        if isinstance(task, ir.Logical):
-            return self.handle_logical(task, ctx)
-        else:
-            return task
-
-    def handle_logical(self, task: ir.Logical, ctx: Context):
-        # Process the original body to find candidates for extraction. Each task is in one
-        # of three categories:
-        # - Binders: tasks that bind variables. These are candidates for extracting into
-        #   the common body.
-        # - Flattenables: tasks that will later be extracted by the Flatten pass
-        # - Other: tasks that are neither binders nor flattenables; these will remain
-        #   in the body as-is.
-
-        binders = ordered_set()
-        flattenables = ordered_set()
-        other = ordered_set()
-
-        for child in task.body:
-            if _is_binder(child):
-                binders.add(child)
-            elif _is_flattenable(ctx, child):
-                flattenables.add(child)
-            else:
-                other.add(child)
-
-        # The new body of the rewritten task
-        body:OrderedSet[ir.Task] = ordered_set()
-
-        # Quick check to see if it's worth doing more analysis; we only want to extract
-        # common binders if there are multiple, and there are also multiple flattenables
-        # that will be extracted by the flatten pass later (so that they can share the
-        # extracted logic).
-        plan: Optional[ExtractCommon.ExtractionPlan] = None
-        if len(binders) > 1 and len(flattenables) > 1:
-            plan = _create_extraction_plan(ctx, binders, flattenables, other)
-            if plan:
-                # plan is worthwhile, extract the common body and add the connection to the body
-                exposed_vars = plan.exposed_vars.get_list()
-                plan.common_reference = f.lookup(helpers.extract(task, plan.common_body, exposed_vars, ctx.rewrite_ctx, "common"), exposed_vars)
-
-                # Add plan common reference to the body.
-                body.add(plan.common_reference)
-
-        # recursively handle children
-        for child in task.body:
-            # skip children that were extracted
-            if plan and child not in other and child not in plan.remaining_body and child not in flattenables:
-                continue
-
-            # no plan or child is not a composite, so just add the handled to the body
-            if not plan or child not in flattenables:
-                body.add(self.handle(child, ctx))
-                continue
-
-            # there is a plan and the child is in composites, so...
-            replacement = self.handle(child, ctx)
-
-            # this child needs either extra local dependencies or the common reference
-            if child in plan.local_dependencies:
-                # the new body will have maybe the common reference and the local deps
-                replacement_body = ordered_set()
-
-                hoisted = OrderedSet()
-                if isinstance(replacement, ir.Logical):
-                    # if replacement is a logical, just keep the same hoisted vars
-                    hoisted.update(replacement.hoisted)
-                else:
-                    # otherwise, we need to hoist the vars that are output from local deps
-                    # and input to the replacement task
-                    dep_outputs = OrderedSet()
-                    for d in plan.local_dependencies.get(child, ordered_set()):
-                        dep_outputs.update(ctx.info.task_outputs(d))
-                    hoisted.update(dep_outputs & ctx.info.task_inputs(replacement))
-
-                if child in plan.local_dependencies:
-                    for local_dep in plan.local_dependencies[child]:
-                        replacement_body.add(local_dep.clone())
-
-                if isinstance(replacement, ir.Logical):
-                    # if the replacements is a logical, we can just add to the body
-                    body.add(replacement.reconstruct(
-                        replacement.engine,
-                        tuple(hoisted.get_list()),
-                        tuple(replacement_body.update(replacement.body).get_list()),
-                        replacement.annotations
-                    ))
-                else:
-                    # Otherwise, wrap the local dependencies in a Lookup where the output
-                    # variables are hoisted, and keep the computed replacement.
-                    body.add(f.logical(replacement_body.get_list(), hoisted.get_list(), replacement.engine))
-                    body.add(replacement)
-            else:
-                # child does not need extras in the body, just add it to the main body
-                body.add(replacement)
-
-        return ir.Logical(task.engine, task.hoisted, tuple(body))
-
-
-#--------------------------------------------------
-# Utilities
-#--------------------------------------------------
-
-def _create_extraction_plan(ctx: ExtractCommon.Context, binders: OrderedSet[ir.Task], flattenables: OrderedSet[ir.Task], others: OrderedSet[ir.Task]) -> Optional[ExtractCommon.ExtractionPlan]:
-    """
-    Compute a plan to extract tasks in this frame that are common dependencies
-    across these composite tasks.
-    """
-    # If there are any pragma lookups, then don't extract anything. Pragma lookups are
-    # designed to control execution order, and extracting them may affect their
-    # semantics.
-    for b in binders:
-        if isinstance(b, ir.Lookup) and builtins.is_pragma(b.relation):
-            return None
-
-    # Compute intersection of task dependencies
-    all_deps = [ctx.info.task_dependencies(f) for f in flattenables]
-    deps = reduce(lambda a, b: a & b, all_deps)
-    common_body = binders & deps
-
-    # We don't need to extract anything if there's only zero or one common tasks
-    if len(common_body) < 2:
-        return None
-
-    # Keep track of remaining tasks that are not extracted in the common body
-    remaining = ordered_set()
-
-    # Compute the vars that should be output from the common body. These are the union of
-    # all input vars across all non-extracted tasks, intersected with output vars of
-    # the common body.
-
-    # First, compute the output vars of the common body
-    common_body_output_vars = OrderedSet()
-    for child in common_body:
-        common_body_output_vars.update(ctx.info.task_outputs(child))
-
-    # Next, compute the union of the input vars of all non-extracted tasks
-    non_extracted_tasks = (binders - common_body) | flattenables | others
-    all_exposed_vars: list[OrderedSet[ir.Var]] = []
-    for t in non_extracted_tasks:
-        input_vars = ctx.info.task_inputs(t)
-        all_exposed_vars.append(input_vars if input_vars else OrderedSet())
-
-    exposed_vars = reduce(lambda a, b: a | b, all_exposed_vars) & common_body_output_vars
-
-    # If there are no vars in common, then it's not worth extracting
-    if not exposed_vars:
-        return None
-
-    # Make sure that all local dependencies of the common body are included in the common
-    # body. This is important for the safety of this rewrite.
-    for task in common_body:
-        local_deps = ctx.info.local_dependencies(task)
-        if local_deps:
-            common_body.update(local_deps & binders)
-
-    # check which of the original binders remain, and make sure their dependencies also stay
-    for binder in binders:
-        if binder not in common_body:
-            remaining.add(binder)
-            deps = _compute_local_dependencies(ctx, binders, binder, exposed_vars)
-            if deps:
-                remaining.update(deps)
-
-    # for each composite, check if there are additional tasks needed, because the task
-    # depends on it but it is not exposed by the vars
-    local_dependencies: dict[ir.Task, OrderedSet[ir.Task]] = dict()
-    for flattenable in flattenables:
-        local = _compute_local_dependencies(ctx, binders, flattenable, exposed_vars)
-        if local:
-            local_dependencies[flattenable] = local
-
-    return ExtractCommon.ExtractionPlan(common_body, remaining, exposed_vars, local_dependencies)
-
-def _compute_local_dependencies(ctx: ExtractCommon.Context, binders: OrderedSet[ir.Task], composite: ir.Task, exposed_vars: OrderedSet[ir.Var]):
-    """
-    The tasks in common_body will be extracted into a logical that will expose the exposed_vars.
-    Compute which additional dependencies are needed specifically for this composite, because
-    it depends on some tasks that are extracted to common_body but not exposed by exposed_vars.
-    """
-
-    # working list of vars we still need to fulfill
-    inputs = ctx.info.task_inputs(composite)
-    if not inputs:
-        return None
-
-    # vars exposed by exposed vars + tasks added to the local body
-    vars_exposed = OrderedSet.from_iterable(exposed_vars)
-    vars_needed = (inputs - vars_exposed)
-    if not vars_needed:
-        return None
-
-    # this is a greedy algorithm that uses the first task in the common body that provides
-    # a variable needed; it may result in sub-optimal extraction, but should be correct
-    local_body = ordered_set()
-    while(vars_needed):
-        v = vars_needed.pop()
-        for x in binders:
-            if x not in local_body:
-                # an x that is not yet in local_body can fulfill v
-                x_outputs = ctx.info.task_outputs(x)
-                if x_outputs and v in x_outputs:
-                    # add it to local_body and add its outputs to vars exposed
-                    local_body.add(x)
-                    vars_exposed.add(x_outputs)
-                    # but add its inputs the vars now needed
-                    inputs = ctx.info.task_inputs(x)
-                    if inputs:
-                        vars_needed.update(inputs - vars_exposed)
-    return local_body
-
-def _is_binder(task: ir.Task):
-    binder_types = (ir.Lookup, ir.Construct, ir.Exists, ir.Data, ir.Not)
-
-    # If the task itself is a binder
-    if isinstance(task, binder_types):
-        return True
-
-    # If the task is a Logical containing only binders
-    if isinstance(task, ir.Logical) and all(_is_binder(c) for c in task.body):
-        return True
-
-    # If the task is a Union containing only binders
-    if isinstance(task, ir.Union) and all(_is_binder(c) for c in task.tasks):
-        return True
-
-    return False
-
-def _is_flattenable(ctx: ExtractCommon.Context, task: ir.Task):
-    # Each output will be flattened into its own top-level def
-    if isinstance(task, ir.Output):
-        return True
-
-    extractable_types = (ir.Update, ir.Aggregate, ir.Match, ir.Rank)
-    return isinstance(task, ir.Logical) and len(visitor.collect_by_type(extractable_types, task)) > 0