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

relationalai/semantics/lqp/rewrite/cdc.py

@@ -0,0 +1,216 @@
+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

@@ -0,0 +1,338 @@
+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