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

relationalai/semantics/lqp/rewrite/cdc.py

@@ -1,216 +0,0 @@
-from __future__ import annotations
-from typing import Optional
-
-from relationalai.semantics.metamodel import ir, factory as f, builtins as bt, types, helpers
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
-from relationalai.semantics.metamodel import dependency
-
-class CDC(Pass):
-    """
-    Pass to process tables brought to Relational AI logical engines by CDC. When CDC occurs,
-    wide snowflake tables are shredded into smaller tables. This pass ensures that code that
-    reads from the wide relation is changed to read from the smaller tables. Furthermore,
-    it attaches the @function annotation to the property lookups, as an optimization.
-
-    Beware that this pass makes assumptions about the names and types of CDC relations and
-    columns!
-
-    From:
-        Logical
-            TPCH.SF1.LINEITEM(l_orderkey, l_partkey, l_suppkey, l_linenumber, l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate, l_commitdate, l_receiptdate, l_shipinstruct, l_shipmode, l_comment)
-            construct(LineItem, "l_orderkey", l_orderkey, "l_linenumber", l_linenumber, lineitem)
-            -> derive LineItem(lineitem)
-            -> derive l_orderkey(lineitem, l_orderkey)
-            -> derive l_linenumber(lineitem, l_linenumber)
-    To:
-    Logical
-        tpch_sf1_lineitem("L_ORDERKEY", row_id, l_orderkey)
-        tpch_sf1_lineitem("L_LINENUMBER", row_id, l_linenumber)
-        construct(LineItem, "l_orderkey", l_orderkey, "l_linenumber", l_linenumber, lineitem)
-        -> derive LineItem(lineitem)
-        -> derive l_orderkey(lineitem, l_orderkey) (@function)
-        -> derive l_linenumber(lineitem, l_linenumber) (@function)
-    """
-
-    #--------------------------------------------------
-    # Public API
-    #--------------------------------------------------
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        # create the dependency analysis context
-        ctx = CDC.CDCContext(model)
-
-        # 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
-        )
-
-    class CDCContext():
-        def __init__(self, model: ir.Model):
-            self.model = model
-            self.rewrite_ctx = helpers.RewriteContext()
-            self.cdc_relations = dict()
-            self.info: Optional[dependency.DependencyInfo] = None
-
-        def get_info(self):
-            """ Lazily compute info as needed. """
-            if self.info is None:
-                # no need for dependency analsyis, only inputs/outputs
-                self.info = dependency.analyze_bindings(self.model.root)
-            return self.info
-
-    #--------------------------------------------------
-    # IR handlers
-    #--------------------------------------------------
-
-    def handle(self, task: ir.Task, ctx: CDC.CDCContext):
-        # 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)
-        elif isinstance(task, ir.Not):
-            return self.handle_not(task, ctx)
-        else:
-            return task
-
-    def handle_logical(self, task: ir.Logical, ctx: CDC.CDCContext):
-
-        wide_cdc_table_lookups = ordered_set()
-        for child in task.body:
-            if isinstance(child, ir.Lookup) and bt.from_cdc_annotation in child.relation.annotations:
-                wide_cdc_table_lookups.add(child)
-
-        # optimization to avoid creating a frame if unnecessary
-        if not wide_cdc_table_lookups:
-            # no need to analyze dependencies, just handle children recursively and
-            # reconstruct the logical
-            body:OrderedSet[ir.Task] = ordered_set()
-            for child in task.body:
-                body.add(self.handle(child, ctx))
-            return ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations)
-
-        # ensure function annotation is in the model
-        # TODO: need to revisit this to add @function annotations only when we really need them
-        # ctx.analysis_ctx.relations.append(rel_bt.function)
-
-        # process the children
-        body:OrderedSet[ir.Task] = ordered_set()
-
-        # find variables required by the other tasks
-        required_vars = ordered_set()
-        for child in task.body:
-            if child not in wide_cdc_table_lookups:
-                required_vars.update(ctx.get_info().task_inputs(child))
-                #
-                # Vars used in atoms (outputs) may be column variables:
-                # e.g.:
-                #       Transaction.type(source.COLUMN)
-                #    or
-                #       Transaction.type = source.COLUMN
-                #
-                # Neither of these mark their vars as inputs, so we pull from the outputs as well.
-                # The only exception is RowId, since there's no need to join on `METADATA$KEY` if
-                # a non-table atom has `RowId` in keys.
-                #
-                output_vars = ctx.get_info().task_outputs(child) or []
-                output_vars = [var for var in output_vars if not types.matches(var.type, types.RowId)]
-                required_vars.update(output_vars)
-
-        # We must return anything that's hoisted
-        required_vars.update(helpers.hoisted_vars(task.hoisted))
-
-        # rewrite the cdc table lookup into lookups for each required variable
-        for child in task.body:
-            if child in wide_cdc_table_lookups:
-                assert isinstance(child, ir.Lookup)
-                wide_relation = child.relation
-                properties = required_vars & ctx.get_info().task_outputs(child)
-                if properties:
-                    assert isinstance(child.args[0], ir.Var) and types.matches(child.args[0].type, types.RowId)
-                    row_id = child.args[0]
-                    for property in properties:
-                        if types.matches(property.type, types.RowId) and len(properties) > 1:
-                            continue
-
-                        relation = self._get_property_cdc_relation(wide_relation, property, ctx)
-                        field_name = ir.Literal(types.Symbol, property.name)
-                        if types.matches(property.type, types.RowId):
-                            field_name = ir.Literal(types.Symbol, "METADATA$KEY")
-                            property = ir.Var(type=types.RowId, name=property.name)
-
-                            # METADATA$KEY is unary
-                            relation = f.relation(
-                                relation.name,
-                                [f.field("symbol", types.Symbol), f.field("row_id", types.RowId)],
-                                annos=[*relation.annotations],
-                            )
-                            body.add(ir.Lookup(
-                                task.engine,
-                                relation,
-                                tuple([field_name, row_id])
-                            ))
-                        else:
-                            body.add(ir.Lookup(
-                                task.engine,
-                                relation,
-                                tuple([field_name, row_id, property])
-                            ))
-
-            # handle non cdc table children, adding @function to the updates
-            for child in task.body:
-                if child not in wide_cdc_table_lookups:
-                    body.add(self.handle(child, ctx))
-
-                    # TODO: need to revisit this to add @function annotations only when we really need them
-                    # replacement = self.handle(child, ctx)
-                    # if isinstance(replacement, ir.Update):
-                    #     if len(replacement.args) == 1:
-                    #         body.add(replacement)
-                    #     else:
-                    #         body.add(replacement.reconstruct(
-                    #             replacement.engine,
-                    #             replacement.relation,
-                    #             replacement.args,
-                    #             replacement.effect,
-                    #             replacement.annotations | [rel_bt.function_annotation]
-                    #         ))
-                    # else:
-                    #     body.add(replacement)
-
-        return ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations)
-
-    def handle_not(self, not_task: ir.Not, ctx: CDC.CDCContext):
-        sub_task = self.handle(not_task.task, ctx)
-        return ir.Not(not_task.engine, sub_task, not_task.annotations)
-
-    def _get_property_cdc_relation(self, wide_cdc_relation: ir.Relation, property: ir.Var, ctx: CDC.CDCContext):
-        """
-        Get the relation that represents this property var in this wide_cdc_relation. If the
-        relation is not yet available in the context, this method will create and register it.
-        """
-        relation_name = helpers.sanitize(wide_cdc_relation.name).replace("-", "_")
-        key = (relation_name, property.name)
-        if key not in ctx.cdc_relations:
-            # the property relation is overloaded for all properties of the same wide cdc relation, so they have
-            # the same name, but potentially a different type in the value column; also note that they are
-            # annotated as external to avoid renaming.
-            relation = f.relation(
-                relation_name,
-                [f.field("symbol", types.Symbol), f.field("row_id", types.Number), f.field("value", property.type)],
-                annos=[bt.external_annotation]
-            )
-            ctx.cdc_relations[key] = relation
-            ctx.rewrite_ctx.relations.append(relation)
-        return ctx.cdc_relations[key]

relationalai/semantics/lqp/rewrite/extract_common.py

@@ -1,338 +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):
-    # If the task itself is a binder
-    if any(isinstance(task, binder) for binder in (ir.Lookup, ir.Construct, ir.Exists, ir.Data, ir.Not)):
-        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