relationalai 0.13.2__py3-none-any.whl → 0.13.3__py3-none-any.whl

relationalai/semantics/lqp/rewrite/eliminate_data.py

@@ -0,0 +1,108 @@
+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/period_math.py

@@ -0,0 +1,77 @@
+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 types
+
+# Generate date arithmetic expressions, such as
+# `rel_primitive_date_add(:day, [date] delta, res_2)` by finding the period
+# expression for the delta and adding the period type to the date arithmetic expression.
+#
+# date_add and it's kin are generated by a period expression, e.g.,
+# `day(delta, res_1)`
+# followed by the date arithmetic expression using the period
+# `date_add([date] res_1 res_2)`
+class PeriodMath(Pass):
+    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
+        period_rewriter = self.PeriodRewriter()
+        model = period_rewriter.walk(model)
+        period_math_rewriter = self.PeriodMathRewriter(period_rewriter.period_vars)
+        model = period_math_rewriter.walk(model)
+        return model
+
+    # Find all period builtins. We need to make them safe for the emitter (either by
+    # translating to a cast, or removing) and store the variable and period type for use
+    # in the date/datetime add/subtract expressions.
+    class PeriodRewriter(visitor.Rewriter):
+        def __init__(self):
+            super().__init__()
+            self.period_vars: dict[ir.Var, str] = {}
+
+        def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
+            if not rel_builtins.is_builtin(node.relation):
+                return node
+
+            if node.relation.name not in {
+                "year", "month", "week", "day", "hour", "minute", "second", "millisecond", "microsecond", "nanosecond"
+            }:
+                return node
+
+            assert len(node.args) == 2, "Expect 2 arguments for period builtins"
+            assert isinstance(node.args[1], ir.Var), "Expect result to be a variable"
+            period = node.relation.name
+            result_var = node.args[1]
+            self.period_vars[result_var] = period
+
+            # Ideally we could now remove the unused and unhandled period type construction
+            # but we may also need to cast the original variable to an Int64 for use by the
+            # date/datetime add/subtract expressions.
+            # TODO: Remove the node entirely where possible and update uses of the result
+            return f.lookup(rel_builtins.cast, [types.Int64, node.args[0], result_var])
+
+    # Update date/datetime add/subtract expressions with period information.
+    class PeriodMathRewriter(visitor.Rewriter):
+        def __init__(self, period_vars: dict[ir.Var, str]):
+            super().__init__()
+            self.period_vars: dict[ir.Var, str] = period_vars
+
+        def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
+            if not rel_builtins.is_builtin(node.relation):
+                return node
+
+            if node.relation.name not in {
+                "date_add", "date_subtract", "datetime_add", "datetime_subtract"
+            }:
+                return node
+
+            if len(node.args) == 4:
+                # We've already visited this lookup
+                return node
+
+            assert isinstance(node.args[1], ir.Var), "Expect period to be a variable"
+            period_var = node.args[1]
+            assert period_var in self.period_vars, "datemath found, but no vars to insert"
+
+            period = self.period_vars[period_var]
+
+            new_args = [f.literal(period, types.Symbol)] + [arg for arg in node.args]
+
+            return f.lookup(node.relation, new_args)

relationalai/semantics/lqp/rewrite/quantify_vars.py

@@ -69,7 +69,7 @@ class VarScopeInfo(Visitor):
                     ir.Var, ir.Literal, ir.Relation, ir.Field,
                     ir.Default, ir.Output, ir.Update, ir.Aggregate,
                     ir.Check, ir.Require,
-                    ir.Annotation, ir.Rank)
+                    ir.Annotation, ir.Rank, ir.Break)
 
     def __init__(self):
         super().__init__()
@@ -103,16 +103,29 @@ class VarScopeInfo(Visitor):
             self._record(node, scope_vars)
 
         elif isinstance(node, (ir.Match, ir.Union)):
-            # Match/Union inherits the vars in scope from its sub-tasks.
+            # Match/Union only inherits vars if they are in scope for all sub-tasks.
             scope_vars = ordered_set()
+            # Prime the search with the first sub-task's vars.
+            if node.tasks:
+                scope_vars.update(self._vars_in_scope.get(node.tasks[0].id, None))
+
             for task in node.tasks:
                 sub_scope_vars = self._vars_in_scope.get(task.id, None)
-                if sub_scope_vars:
-                    scope_vars.update(sub_scope_vars)
+                if not scope_vars or not sub_scope_vars:
+                    scope_vars = ordered_set()
+                    break
+                scope_vars = (scope_vars & sub_scope_vars)
+
             # Hoisted vars are not considered for quantification at this level.
             scope_vars.difference_update(helpers.hoisted_vars(node.hoisted))
             self._record(node, scope_vars)
 
+        elif isinstance(node, (ir.Loop, ir.Sequence)):
+            # Variables in Loops and Sequences are scoped exclusively within the body and
+            # not propagated outside. No need to record any variables, as they shouldn't be
+            # in scope for the node itself
+            pass
+
         elif isinstance(node, ir.Logical):
             self._do_logical(node)
 
@@ -128,6 +141,9 @@ class VarScopeInfo(Visitor):
         all_nested_vars = ordered_set()
         output_vars = ordered_set()
 
+        # Collect variables nested in child Logical and Not nodes
+        nested_vars_in_task: dict[ir.Var, int] = dict()
+
         # Collect all variables from logical sub-tasks
         for task in node.body:
             if isinstance(task, ir.Output):
@@ -140,19 +156,29 @@ class VarScopeInfo(Visitor):
                         scope_vars.add(var)
                 continue
 
-            sub_scope_vars = self._vars_in_scope.get(task.id, None)
-
             # Hoisted variables from sub-tasks are brought again into scope.
             if isinstance(task, (ir.Logical, ir.Union, ir.Match)):
                 scope_vars.update(helpers.hoisted_vars(task.hoisted))
 
-            if sub_scope_vars:
-                if isinstance(task, ir.Logical):
-                    all_nested_vars.update(sub_scope_vars)
-                elif not isinstance(task, ir.Not):
-                    # For all other node kinds (except Not), just propagate the variables in scope.
-                    # Not nodes stop the propagation of variables coming from their sub-tasks.
-                    scope_vars.update(sub_scope_vars)
+            # Get variables in sub-task scope
+            sub_scope_vars = self._vars_in_scope.get(task.id, ordered_set())
+
+            if isinstance(task, ir.Logical):
+                # Logical child nodes should have their nested variables quantified
+                # only if they are needed in more than one child task.
+                for var in sub_scope_vars:
+                    if var not in nested_vars_in_task:
+                        nested_vars_in_task[var] = 0
+                    nested_vars_in_task[var] += 1
+            elif not isinstance(task, ir.Not):
+                # Other nodes with nested variables need to be quantified at this level
+                scope_vars.update(sub_scope_vars)
+
+        for v, c in nested_vars_in_task.items():
+            # If the variable appears in more than one nested child, then it needs to be
+            # quantified here. Otherwise, it will be handled in the child node
+            if c > 1:
+                all_nested_vars.add(v)
 
         # Nested variables also need to be introduced, provided they are not output variables.
         for var in all_nested_vars:
@@ -190,37 +216,30 @@ class FindQuantificationNodes(Visitor):
     def __init__(self, var_info: VarScopeInfo):
         super().__init__()
         self._vars_in_scope = var_info._vars_in_scope
-        self._handled_vars = ordered_set()
+        self.handled_vars: dict[int, OrderedSet[ir.Var]] = {}
         self.node_quantifies_vars = {}
 
     def enter(self, node: ir.Node, parent: Optional[ir.Node]=None) -> "Visitor":
         if contains_only_declarable_constraints(node):
             return self
 
+        handled_vars = self.handled_vars.get(parent.id, ordered_set()) if parent else ordered_set()
+        # Clone the set to avoid modifying parent's handled vars
+        handled_vars = OrderedSet.from_iterable(handled_vars)
+
         if isinstance(node, (ir.Logical, ir.Not)):
             ignored_vars = _ignored_vars(node)
-            self._handled_vars.update(ignored_vars)
+            handled_vars.update(ignored_vars)
 
             scope_vars = self._vars_in_scope.get(node.id, None)
             if scope_vars:
-                scope_vars.difference_update(self._handled_vars)
+                scope_vars.difference_update(handled_vars)
                 if scope_vars:
-                    self._handled_vars.update(scope_vars)
+                    handled_vars.update(scope_vars)
                     self.node_quantifies_vars[node.id] = scope_vars
-        return self
-
-    def leave(self, node: ir.Node, parent: Optional[ir.Node]=None) -> ir.Node:
-        if contains_only_declarable_constraints(node):
-            return node
 
-        if isinstance(node, (ir.Logical, ir.Not)):
-            ignored_vars = _ignored_vars(node)
-            self._handled_vars.difference_update(ignored_vars)
-
-            scope_vars = self._vars_in_scope.get(node.id, None)
-            if scope_vars:
-                self._handled_vars.difference_update(scope_vars)
-        return node
+        self.handled_vars[node.id] = handled_vars
+        return self
 
 class QuantifyVarsRewriter(Rewriter):
     """
@@ -254,7 +273,12 @@ class QuantifyVarsRewriter(Rewriter):
                     # in IR directly may do so and the flatten pass doesn't split them yet.
                     if len(agg_or_rank_tasks) > 0:
                         print(f"Multiple aggregate/rank tasks found: {agg_or_rank_tasks} and {task}")
-                    agg_or_rank_tasks.append(task)
+                    # If the agg/rank depends on any of the vars being quantified here,
+                    # then it needs to be inside the quantification
+                    if any(var in helpers.vars(task.projection) for var in vars):
+                        inner_tasks.append(task)
+                    else:
+                        agg_or_rank_tasks.append(task)
 
                 else:
                     inner_tasks.append(task)
@@ -283,6 +307,16 @@ class QuantifyVarsRewriter(Rewriter):
 
         return node if node.task is new_task else f.not_(new_task)
 
+    def handle_union(self, node: ir.Union, parent: ir.Node, ctx:Optional[Any]=None) -> ir.Union:
+        if not node.tasks:
+            return node
+
+        new_tasks = self.walk_list(node.tasks, node)
+        return node if node.tasks is new_tasks else f.union(
+            tasks = new_tasks,
+            hoisted = node.hoisted,
+        )
+
     # To avoid unnecessary cloning of vars in the visitor.
     def handle_var(self, node: ir.Var, parent: ir.Node, ctx:Optional[Any]=None) -> ir.Var:
         return node

relationalai/semantics/lqp/rewrite/unify_definitions.py

@@ -0,0 +1,317 @@
+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 relationalai.semantics.metamodel import helpers
+from relationalai.semantics.metamodel.util import FrozenOrderedSet, OrderedSet
+
+
+from typing import cast, Union, Optional, Iterable
+from collections import defaultdict
+
+# LQP does not support multiple definitions for the same relation. This pass unifies all
+# definitions for each relation into a single definition using a union.
+class UnifyDefinitions(Pass):
+    def __init__(self):
+        super().__init__()
+
+    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
+        # Maintain a cache of renamings for each relation. These need to be consistent
+        # across all definitions of the same relation.
+        self.renamed_relation_args: dict[Union[ir.Value, ir.Relation], list[ir.Var]] = {}
+
+        root = cast(ir.Logical, model.root)
+        new_tasks = self.get_combined_multidefs(root)
+        return ir.Model(
+            model.engines,
+            model.relations,
+            model.types,
+            f.logical(
+                tuple(new_tasks),
+                root.hoisted,
+                root.engine,
+            ),
+            model.annotations,
+        )
+
+    def _get_heads(self, logical: ir.Logical) -> list[Union[ir.Update, ir.Output]]:
+        derives = []
+        for task in logical.body:
+            if isinstance(task, ir.Update) and task.effect == ir.Effect.derive:
+                derives.append(task)
+            elif isinstance(task, ir.Output):
+                derives.append(task)
+        return derives
+
+    def _get_non_heads(self, logical: ir.Logical) -> list[ir.Task]:
+        non_derives = []
+        for task in logical.body:
+            if not(isinstance(task, ir.Update) and task.effect == ir.Effect.derive) and not isinstance(task, ir.Output):
+                non_derives.append(task)
+        return non_derives
+
+    def _get_head_identifier(self, head: Union[ir.Update, ir.Output]) -> Optional[ir.Value]:
+        if isinstance(head, ir.Update):
+            return head.relation
+        else:
+            assert isinstance(head, ir.Output)
+            if len(head.aliases) <= 2:
+                # For processing here, we need output to have at least the column markers
+                # `cols` and `col`, and also a key
+                return None
+
+            output_alias_names = helpers.output_alias_names(head.aliases)
+            output_vals = helpers.output_values(head.aliases)
+
+            # For normal outputs, the pattern is output[keys](cols, "col000" as 'col', ...)
+            if output_alias_names[0] == "cols" and output_alias_names[1] == "col":
+                return output_vals[1]
+
+            # For exports, the pattern is output[keys]("col000" as 'col', ...)
+            if helpers.is_export(head):
+                if output_alias_names[0] == "col":
+                    return output_vals[0]
+
+        return None
+
+    def get_combined_multidefs(self, root: ir.Logical) -> list[ir.Logical]:
+        # Step 1: Group tasks by the relation they define.
+        relation_to_tasks: dict[Union[None, ir.Value, ir.Relation], list[ir.Logical]] = defaultdict(list)
+
+        for task in root.body:
+            task = cast(ir.Logical, task)
+            task_heads = self._get_heads(task)
+
+            # Some relations do not need to be grouped, e.g., if they don't contain a
+            # derive. Use `None` as a placeholder key for these cases.
+            if len(task_heads) != 1:
+                relation_to_tasks[None].append(task)
+                continue
+
+            head_id = self._get_head_identifier(task_heads[0])
+            relation_to_tasks[head_id].append(task)
+
+        # Step 2: For each relation, combine all of the body definitions into a union.
+        result_tasks = []
+        for relation, tasks in relation_to_tasks.items():
+            # If there's only one task for the relation, or if grouping is not needed, then
+            # just keep the original tasks.
+            if len(tasks) == 1 or relation is None:
+                result_tasks.extend(tasks)
+                continue
+
+            result_tasks.append(self._combine_tasks_into_union(tasks))
+        return result_tasks
+
+    def _get_variable_mapping(self, logical: ir.Logical) -> dict[ir.Value, ir.Var]:
+        heads = self._get_heads(logical)
+        assert len(heads) == 1, "should only have one head in a logical at this stage"
+        head = heads[0]
+
+        var_mapping = {}
+        head_id = self._get_head_identifier(head)
+
+        if isinstance(head, ir.Update):
+            args_for_renaming = head.args
+        else:
+            assert isinstance(head, ir.Output)
+            output_alias_names = helpers.output_alias_names(head.aliases)
+            if output_alias_names[0] == "cols" and output_alias_names[1] == "col":
+                assert len(head.aliases) > 2
+
+                # For outputs, we do not need to rename the `cols` and `col` markers or the
+                # keys.
+                output_values = helpers.output_values(head.aliases)[2:]
+
+            else:
+                assert helpers.is_export(head) and output_alias_names[0] == "col"
+                assert len(head.aliases) > 1
+
+                # For exports, we do not need to rename the `col` marker or the keys.
+                output_values = helpers.output_values(head.aliases)[1:]
+
+            args_for_renaming = []
+            for v in output_values:
+                if head.keys and isinstance(v, ir.Var) and v in head.keys:
+                    continue
+                args_for_renaming.append(v)
+
+        if head_id not in self.renamed_relation_args:
+            renamed_vars = []
+            for (i, arg) in enumerate(args_for_renaming):
+                typ = typer.to_type(arg)
+                assert arg not in var_mapping, "args of update should be unique"
+                if isinstance(arg, ir.Var):
+                    var_mapping[arg] = ir.Var(typ, arg.name)
+                else:
+                    var_mapping[arg] = ir.Var(typ, f"arg_{i}")
+
+                renamed_vars.append(var_mapping[arg])
+            self.renamed_relation_args[head_id] = renamed_vars
+        else:
+            for (arg, var) in zip(args_for_renaming, self.renamed_relation_args[head_id]):
+                var_mapping[arg] = var
+
+        return var_mapping
+
+    def _rename_variables(self, logical: ir.Logical) -> ir.Logical:
+        class RenameVisitor(visitor.Rewriter):
+            def __init__(self, var_mapping: dict[ir.Value, ir.Var]):
+                super().__init__()
+                self.var_mapping = var_mapping
+
+            def _get_mapped_value(self, val: ir.Value) -> ir.Value:
+                if isinstance(val, tuple):
+                    return tuple(self._get_mapped_value(t) for t in val)
+                return self.var_mapping.get(val, val)
+
+            def _get_mapped_values(self, vals: Iterable[ir.Value]) -> list[ir.Value]:
+                return [self._get_mapped_value(v) for v in vals]
+
+            def handle_var(self, node: ir.Var, parent: ir.Node) -> ir.Var:
+                return self.var_mapping.get(node, node)
+
+            # TODO: ideally, extend the rewriter class to allow rewriting PyValue to Var so
+            # we don't need to separately handle all cases containing them.
+            def handle_update(self, node: ir.Update, parent: ir.Node) -> ir.Update:
+                return ir.Update(
+                    node.engine,
+                    node.relation,
+                    tuple(self._get_mapped_values(node.args)),
+                    node.effect,
+                    node.annotations,
+                )
+
+            def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
+                return ir.Lookup(
+                    node.engine,
+                    node.relation,
+                    tuple(self._get_mapped_values(node.args)),
+                    node.annotations,
+                )
+
+            def handle_output(self, node: ir.Output, parent: ir.Node) -> ir.Output:
+                new_aliases = FrozenOrderedSet(
+                    [(name, self._get_mapped_value(value)) for name, value in node.aliases]
+                )
+                if node.keys:
+                    new_keys = FrozenOrderedSet(
+                        [self.var_mapping.get(key, key) for key in node.keys]
+                    )
+                else:
+                    new_keys = node.keys
+
+                return ir.Output(
+                    node.engine,
+                    new_aliases,
+                    new_keys,
+                    node.annotations,
+                )
+
+            def handle_construct(self, node: ir.Construct, parent: ir.Node) -> ir.Construct:
+                new_values = tuple(self._get_mapped_values(node.values))
+                new_id_var = self.var_mapping.get(node.id_var, node.id_var)
+                return ir.Construct(
+                    node.engine,
+                    new_values,
+                    new_id_var,
+                    node.annotations,
+                )
+
+            def handle_aggregate(self, node: ir.Aggregate, parent: ir.Node) -> ir.Aggregate:
+                new_projection = tuple(self.var_mapping.get(arg, arg) for arg in node.projection)
+                new_group = tuple(self.var_mapping.get(arg, arg) for arg in node.group)
+                new_args = tuple(self._get_mapped_values(node.args))
+                return ir.Aggregate(
+                    node.engine,
+                    node.aggregation,
+                    new_projection,
+                    new_group,
+                    new_args,
+                    node.annotations,
+                )
+
+            def handle_rank(self, node: ir.Rank, parent: ir.Node) -> ir.Rank:
+                new_projection = tuple(self.var_mapping.get(arg, arg) for arg in node.projection)
+                new_group = tuple(self.var_mapping.get(arg, arg) for arg in node.group)
+                new_args = tuple(self.var_mapping.get(arg, arg) for arg in node.args)
+                new_result = self.var_mapping.get(node.result, node.result)
+
+                return ir.Rank(
+                    node.engine,
+                    new_projection,
+                    new_group,
+                    new_args,
+                    node.arg_is_ascending,
+                    new_result,
+                    node.limit,
+                    node.annotations,
+                )
+
+        var_mapping = self._get_variable_mapping(logical)
+
+        renamer = RenameVisitor(var_mapping)
+        result = renamer.walk(logical)
+
+        # Also need to append the equality for each renamed constant. E.g., if the mapping
+        # contains (50.0::FLOAT -> arg_2::FLOAT), we need to add
+        # `eq(arg_2::FLOAT, 50.0::FLOAT)` to the result.
+        value_eqs = []
+        for (old_var, new_var) in var_mapping.items():
+            if not isinstance(old_var, ir.Var):
+                value_eqs.append(f.lookup(rel_builtins.eq, [new_var, old_var]))
+
+        return ir.Logical(
+            result.engine,
+            result.hoisted,
+            tuple(value_eqs) + tuple(result.body),
+            result.annotations,
+        )
+
+    # This function is the main workhorse for this rewrite pass. It takes a list of tasks
+    # that define the same relation, and combines them into a single task that defines
+    # the relation using a union of all of the bodies.
+    def _combine_tasks_into_union(self, tasks: list[ir.Logical]) -> ir.Logical:
+        # Step 1: Rename the variables in all tasks so that they will match the final derive
+        # after reconstructing into a union
+        renamed_tasks = [self._rename_variables(task) for task in tasks]
+
+        # Step 2: Get the final derive
+        derives = self._get_heads(renamed_tasks[0])
+        assert len(derives) == 1, "should only have one derive in a logical at this stage"
+        # Also make sure that all the derives are the same. This should be the case because
+        # we renamed all the variables to be the same in step 1.
+        for task in renamed_tasks[1:]:
+            assert self._get_heads(task) == derives, "all derives should be the same"
+
+        derive = derives[0]
+
+        # Step 3: Remove the final `derive` from each task
+        renamed_task_bodies = [
+            f.logical(
+                tuple(self._get_non_heads(t)),  # Only keep non-head tasks
+                t.hoisted,
+                t.engine,
+            )
+            for t in renamed_tasks
+        ]
+
+        # Deduplicate bodies
+        renamed_task_bodies = OrderedSet.from_iterable(renamed_task_bodies).get_list()
+
+        # Step 4: Construct a union of all the task bodies
+        if len(renamed_task_bodies) == 1:
+            # If there's only one body after deduplication, no need to create a union
+            new_body = renamed_task_bodies[0]
+        else:
+            new_body = f.union(
+                tuple(renamed_task_bodies),
+                [],
+                renamed_tasks[0].engine,
+            )
+
+        # Step 5: Add the final derive back
+        return f.logical(
+            (new_body, derive),
+            [],
+            renamed_tasks[0].engine,
+        )