Adjusting how hybrid backrefs are derived to prevent bugs with partitions (#429)

Adjusting how back references are converted from QDAG to Hybrid. These changes eliminate the recursive procedure which utilizes the `back_levels` from the QDAG expression, and replaces it with an iterative solution that only uses the `ancestral_mapping` dictionary, without converting hybrid nodes back into faux QDAG nodes just so they can be recursively re-converted. This required several qdag->relational tests to be adjusted since the qdag structures of those tests were malformed (contained back references to names that had not been pinned with a calculate). Also, after adding a HybridPartition to the pipeline of a HybridTree, appends a Noop afterwards. This ensures that there is a buffer in the `max_steps` field between the child being partitioned by a `HybridPartition` and any other children.

Author: knassre-bodo

pydough/conversion/hybrid_decorrelater.py

@@ -67,7 +67,10 @@ def make_decorrelate_parent(
             to the PARTITION edge case. The third entry is how many operators
             in the pipeline were copied over from the root.
         """
-        if isinstance(hybrid.pipeline[0], HybridPartition) and child_idx == 0:
+        if (
+            isinstance(hybrid.pipeline[0], HybridPartition)
+            and hybrid.children[child_idx].max_steps == 1
+        ):
             # Special case: if the correlated child is the data argument of a
             # partition operation, then the parent to snapshot is actually the
             # parent of the level containing the partition operation. In this

pydough/conversion/hybrid_translator.py

@@ -66,6 +66,7 @@
     HybridCollectionAccess,
     HybridFilter,
     HybridLimit,
+    HybridNoop,
     HybridOperation,
     HybridPartition,
     HybridPartitionChild,
@@ -421,6 +422,7 @@ def populate_children(
                 if (
                     name in hybrid.ancestral_mapping
                     or name in hybrid.pipeline[-1].terms
+                    or subtree.ancestral_mapping[name] == 0
                 ):
                     continue
                 hybrid_back_expr = self.make_hybrid_expr(
@@ -867,97 +869,6 @@ def rewrite_quantile_call(
 
         return max_call
 
-    def make_hybrid_correl_expr(
-        self,
-        back_expr: BackReferenceExpression,
-        collection: PyDoughCollectionQDAG,
-        steps_taken_so_far: int,
-        down_shift: int,
-    ) -> HybridCorrelExpr:
-        """
-        Converts a BACK reference into a correlated reference when the number
-        of BACK levels exceeds the height of the current subtree.
-
-        Args:
-            `back_expr`: the original BACK reference to be converted.
-            `collection`: the collection at the top of the current subtree,
-            before we have run out of BACK levels to step up out of.
-            `steps_taken_so_far`: the number of steps already taken to step
-            up from the BACK node. This is needed so we know how many steps
-            still need to be taken upward once we have stepped out of the child
-            subtree back into the parent subtree.
-            `down_shift`: a factor that should be subtracted from the final
-            back shift when creating a term in the form CORREL(BACK(n).x), to
-            account for edge cases involving PARTITION nodes. Starts as 0, and
-            is incremented as-needed.
-        """
-        if len(self.stack) == 0:
-            raise ValueError("Back reference steps too far back")
-        # Identify the parent subtree that the BACK reference is stepping back
-        # into, out of the child.
-        parent_tree: HybridTree = self.stack.pop()
-        remaining_steps_back: int = back_expr.back_levels - steps_taken_so_far - 1
-        parent_result: HybridExpr
-        new_expr: PyDoughExpressionQDAG
-        # Special case: stepping out of the data argument of PARTITION back
-        # into its ancestor. For example:
-        # TPCH.CALCULATE(x=...).PARTITION(data.WHERE(y > BACK(1).x), ...)
-        partition_edge_case: bool = len(parent_tree.pipeline) == 1 and isinstance(
-            parent_tree.pipeline[0], HybridPartition
-        )
-        if partition_edge_case:
-            next_hybrid: HybridTree
-            if parent_tree.parent is not None:
-                # If the parent tree has a parent, then we can step back
-                # into the parent tree's parent, which is the context for
-                # the partition.
-                next_hybrid = parent_tree.parent
-            else:
-                assert len(self.stack) > 0, "Back reference steps too far back"
-                next_hybrid = self.stack[-1]
-            # Treat the partition's parent as the context for the back
-            # to step into, as opposed to the partition itself (so the back
-            # levels are consistent)
-            self.stack.append(next_hybrid)
-            parent_result = self.make_hybrid_correl_expr(
-                back_expr, collection, steps_taken_so_far, down_shift + 1
-            ).expr
-            self.stack.pop()
-        elif remaining_steps_back == 0:
-            # If there are no more steps back to be made, then the correlated
-            # reference is to a reference from the current context.
-            if back_expr.term_name in parent_tree.ancestral_mapping:
-                new_expr = BackReferenceExpression(
-                    collection,
-                    back_expr.term_name,
-                    parent_tree.ancestral_mapping[back_expr.term_name] - down_shift,
-                )
-                parent_result = self.make_hybrid_expr(parent_tree, new_expr, {}, False)
-            elif back_expr.term_name in parent_tree.pipeline[-1].terms:
-                parent_name: str = parent_tree.pipeline[-1].renamings.get(
-                    back_expr.term_name, back_expr.term_name
-                )
-                parent_result = HybridRefExpr(parent_name, back_expr.pydough_type)
-            else:
-                raise ValueError(
-                    f"Back reference to {back_expr.term_name} not found in parent"
-                )
-        else:
-            # Otherwise, a back reference needs to be made from the current
-            # collection a number of steps back based on how many steps still
-            # need to be taken, and it must be recursively converted to a
-            # hybrid expression that gets wrapped in a correlated reference.
-            new_expr = BackReferenceExpression(
-                collection, back_expr.term_name, remaining_steps_back
-            )
-            parent_result = self.make_hybrid_expr(parent_tree, new_expr, {}, False)
-        # Restore parent_tree back onto the stack, since evaluating `back_expr`
-        # does not change the program's current placement in the subtrees.
-        self.stack.append(parent_tree)
-        # Create the correlated reference to the expression with regards to
-        # the parent tree, which could also be a correlated expression.
-        return HybridCorrelExpr(parent_result)
-
     def add_unique_terms(
         self,
         hybrid: HybridTree,
@@ -1046,6 +957,74 @@ def add_unique_terms(
                 child_idx,
             )
 
+    def translate_back_reference(
+        self, hybrid: HybridTree, expr: BackReferenceExpression
+    ) -> HybridExpr:
+        """
+        Perform the logic used to translate a BACK reference in QDAG into a
+        back reference in hybrid, or a correlated reference if the back
+        reference steps back further than the height of the current hybrid
+        tree.
+
+        Args:
+            `hybrid`: the hybrid tree that should be used to derive the
+            translation of `expr`, as it is the context in which the `expr`
+            will live.
+            `expr`: the BACK reference to be converted.
+
+        Returns:
+            The HybridExpr node corresponding to `expr`.
+        """
+        back_levels: int = 0
+        correl_levels: int = 0
+        new_stack: list[HybridTree] = []
+        ancestor_tree: HybridTree = hybrid
+        expr_name: str = expr.term_name
+        # Start with the current context and hunt for an ancestor with
+        # that name pinned by a CALCULATE (so it is in the ancestral
+        # mapping), and make sure it is within the height bounds of the
+        # current tree. If not, then pop the previous tree from the
+        # stack and look there, repeating until one is found or the
+        # stack is exhausted. Keep track of how many times we step
+        # outward, since this is how many CORREL() layers we need to
+        # wrap the final expression in.
+        while True:
+            if (
+                expr.term_name in ancestor_tree.ancestral_mapping
+                and ancestor_tree.ancestral_mapping[expr.term_name]
+                < ancestor_tree.get_tree_height()
+            ):
+                back_levels = ancestor_tree.ancestral_mapping[expr.term_name]
+                for _ in range(back_levels):
+                    assert ancestor_tree.parent is not None
+                    ancestor_tree = ancestor_tree.parent
+                expr_name = ancestor_tree.pipeline[-1].renamings.get(
+                    expr_name, expr_name
+                )
+                break
+            elif len(self.stack) > 0:
+                ancestor_tree = self.stack.pop()
+                new_stack.append(ancestor_tree)
+                correl_levels += 1
+            else:
+                raise ValueError("Cannot find ancestor with name " + str(expr))
+        for tree in reversed(new_stack):
+            self.stack.append(tree)
+
+        # The final expression is a regular or back reference depending
+        # on how many back levels it is from the identified ancestor.
+        result: HybridExpr
+        if back_levels == 0:
+            result = HybridRefExpr(expr_name, expr.pydough_type)
+        else:
+            result = HybridBackRefExpr(expr_name, back_levels, expr.pydough_type)
+
+        # Then, wrap it in the necessary number of CORREL() layers.
+        for _ in range(correl_levels):
+            result = HybridCorrelExpr(result)
+
+        return result
+
     def make_hybrid_expr(
         self,
         hybrid: HybridTree,
@@ -1074,7 +1053,6 @@ def make_hybrid_expr(
         child_connection: HybridConnection
         args: list[HybridExpr] = []
         hybrid_arg: HybridExpr
-        ancestor_tree: HybridTree
         collection: PyDoughCollectionQDAG
         match expr:
             case PartitionKey():
@@ -1105,30 +1083,8 @@ def make_hybrid_expr(
                 else:
                     return HybridRefExpr(expr.term_name, expr.pydough_type)
             case BackReferenceExpression():
-                # A reference to an expression from an ancestor becomes a
-                # reference to one of the terms of a parent level of the hybrid
-                # tree. If the BACK goes far enough that it must step outside
-                # a child subtree into the parent, a correlated reference is
-                # created.
-                ancestor_tree = hybrid
-                true_steps_back: int = 0
-                # Keep stepping backward until `expr.back_levels` non-hidden
-                # steps have been taken.
-                collection = expr.collection
-                while true_steps_back < expr.back_levels:
-                    assert collection.ancestor_context is not None
-                    collection = collection.ancestor_context
-                    if ancestor_tree.parent is None:
-                        return self.make_hybrid_correl_expr(
-                            expr, collection, true_steps_back, 0
-                        )
-                    ancestor_tree = ancestor_tree.parent
-                    if not ancestor_tree.is_hidden_level:
-                        true_steps_back += 1
-                expr_name = ancestor_tree.pipeline[-1].renamings.get(
-                    expr.term_name, expr.term_name
-                )
-                return HybridBackRefExpr(expr_name, expr.back_levels, expr.pydough_type)
+                return self.translate_back_reference(hybrid, expr)
+
             case Reference():
                 if hybrid.ancestral_mapping.get(expr.term_name, 0) > 0:
                     collection = expr.collection
@@ -1497,6 +1453,8 @@ def make_hybrid_tree(
                         hybrid.pipeline[-1].orderings,
                     )
                 )
+                for name in new_expressions:
+                    hybrid.ancestral_mapping[name] = 0
                 return hybrid
             case Singular():
                 # a Singular node is just used to annotate the preceding context
@@ -1551,6 +1509,11 @@ def make_hybrid_tree(
                 ].subtree
                 partition_child.agg_keys = key_exprs
                 partition_child.join_keys = [(k, k) for k in key_exprs]
+                # Add a dummy no-op after the partition to ensure a
+                # buffer in the max_steps for other children.
+                successor_hybrid.add_operation(
+                    HybridNoop(successor_hybrid.pipeline[-1])
+                )
                 return successor_hybrid
             case OrderBy() | TopK():
                 hybrid = self.make_hybrid_tree(
@@ -1630,6 +1593,11 @@ def make_hybrid_tree(
                         successor_hybrid.children[
                             partition_child_idx
                         ].subtree.agg_keys = key_exprs
+                        # Add a dummy no-op after the partition to ensure a
+                        # buffer in the max_steps for other children.
+                        successor_hybrid.add_operation(
+                            HybridNoop(successor_hybrid.pipeline[-1])
+                        )
                     case GlobalContext():
                         # This is a special case where the child access
                         # is a global context, which means that the child is

pydough/conversion/hybrid_tree.py

@@ -244,6 +244,14 @@ def general_join_condition(self, condition: HybridExpr | None) -> None:
         """
         self._general_join_condition = condition
 
+    def get_tree_height(self) -> int:
+        """
+        Returns the number of levels in the hybrid tree, starting from the
+        current level and counting upward.
+        """
+        parent_height: int = 0 if self.parent is None else self.parent.get_tree_height()
+        return parent_height + 1
+
     def add_operation(self, operation: HybridOperation) -> None:
         """
         Appends a new hybrid operation to the end of the hybrid tree's pipeline.

pydough/conversion/relational_converter.py

@@ -750,7 +750,10 @@ def handle_children(
             # If handling the data for a partition, pull every aggregation key
             # into the current context since it is now accessible as a normal
             # ref instead of a child ref.
-            if isinstance(hybrid.pipeline[pipeline_idx], HybridPartition):
+            if (
+                isinstance(hybrid.pipeline[pipeline_idx], HybridPartition)
+                and child_idx == 0
+            ):
                 partition = hybrid.pipeline[pipeline_idx]
                 assert isinstance(partition, HybridPartition)
                 for key_name in partition.key_names:

tests/test_plan_refsols/alternative_quarter_cum_ir_analysis.txt

@@ -0,0 +1,33 @@
+ROOT(columns=[('quarter', quarter), ('n_incidents', DEFAULT_TO(ndistinct_in_device_id, 0:numeric)), ('n_sold', DEFAULT_TO(n_rows, 0:numeric)), ('quarter_cum', ROUND(RELSUM(args=[DEFAULT_TO(ndistinct_in_device_id, 0:numeric)], partition=[], order=[(quarter):asc_last], cumulative=True) / RELSUM(args=[DEFAULT_TO(n_rows, 0:numeric)], partition=[], order=[(quarter):asc_last], cumulative=True), 2:numeric))], orderings=[(quarter):asc_first])
+ JOIN(condition=t0.quarter == t1.quarter, type=LEFT, cardinality=SINGULAR_FILTER, columns={'n_rows': t1.n_rows, 'ndistinct_in_device_id': t0.ndistinct_in_device_id, 'quarter': t0.quarter})
+  JOIN(condition=t0.quarter == t1.quarter, type=LEFT, cardinality=SINGULAR_FILTER, columns={'ndistinct_in_device_id': t1.ndistinct_in_device_id, 'quarter': t0.quarter})
+   AGGREGATE(keys={'quarter': DATETIME(ca_dt, 'start of quarter':string)}, aggregations={})
+    JOIN(condition=t0.ca_dt < DATETIME(t1.pr_release, '+2 years':string, 'start of quarter':string) & t0.ca_dt >= t1.pr_release, type=INNER, cardinality=SINGULAR_FILTER, columns={'ca_dt': t0.ca_dt})
+     SCAN(table=main.CALENDAR, columns={'ca_dt': ca_dt})
+     FILTER(condition=pr_name == 'RubyCopper-Star':string, columns={'pr_release': pr_release})
+      SCAN(table=main.PRODUCTS, columns={'pr_name': pr_name, 'pr_release': pr_release})
+   AGGREGATE(keys={'quarter': DATETIME(ca_dt, 'start of quarter':string)}, aggregations={'ndistinct_in_device_id': NDISTINCT(in_device_id)})
+    JOIN(condition=t0.in_device_id == t1.de_id & t1.de_product_id == t0.pr_id, type=INNER, cardinality=SINGULAR_FILTER, columns={'ca_dt': t0.ca_dt, 'in_device_id': t0.in_device_id})
+     JOIN(condition=t0.in_repair_country_id == t1.co_id, type=INNER, cardinality=SINGULAR_FILTER, columns={'ca_dt': t0.ca_dt, 'in_device_id': t0.in_device_id, 'pr_id': t1.pr_id})
+      JOIN(condition=t0.ca_dt == DATETIME(t1.in_error_report_ts, 'start of day':string), type=INNER, cardinality=PLURAL_FILTER, columns={'ca_dt': t0.ca_dt, 'in_device_id': t1.in_device_id, 'in_repair_country_id': t1.in_repair_country_id})
+       JOIN(condition=t0.ca_dt < DATETIME(t1.pr_release, '+2 years':string, 'start of quarter':string) & t0.ca_dt >= t1.pr_release, type=INNER, cardinality=SINGULAR_FILTER, columns={'ca_dt': t0.ca_dt})
+        SCAN(table=main.CALENDAR, columns={'ca_dt': ca_dt})
+        FILTER(condition=pr_name == 'RubyCopper-Star':string, columns={'pr_release': pr_release})
+         SCAN(table=main.PRODUCTS, columns={'pr_name': pr_name, 'pr_release': pr_release})
+       SCAN(table=main.INCIDENTS, columns={'in_device_id': in_device_id, 'in_error_report_ts': in_error_report_ts, 'in_repair_country_id': in_repair_country_id})
+      JOIN(condition=True:bool, type=INNER, cardinality=SINGULAR_FILTER, columns={'co_id': t1.co_id, 'pr_id': t0.pr_id})
+       FILTER(condition=pr_name == 'RubyCopper-Star':string, columns={'pr_id': pr_id})
+        SCAN(table=main.PRODUCTS, columns={'pr_id': pr_id, 'pr_name': pr_name})
+       FILTER(condition=co_name == 'CN':string, columns={'co_id': co_id})
+        SCAN(table=main.COUNTRIES, columns={'co_id': co_id, 'co_name': co_name})
+     SCAN(table=main.DEVICES, columns={'de_id': de_id, 'de_product_id': de_product_id})
+  AGGREGATE(keys={'quarter': DATETIME(ca_dt, 'start of quarter':string)}, aggregations={'n_rows': COUNT()})
+   JOIN(condition=t0.de_product_id == t1.pr_id, type=INNER, cardinality=SINGULAR_FILTER, columns={'ca_dt': t0.ca_dt})
+    JOIN(condition=t0.ca_dt == DATETIME(t1.de_purchase_ts, 'start of day':string), type=INNER, cardinality=PLURAL_FILTER, columns={'ca_dt': t0.ca_dt, 'de_product_id': t1.de_product_id})
+     JOIN(condition=t0.ca_dt < DATETIME(t1.pr_release, '+2 years':string, 'start of quarter':string) & t0.ca_dt >= t1.pr_release, type=INNER, cardinality=SINGULAR_FILTER, columns={'ca_dt': t0.ca_dt})
+      SCAN(table=main.CALENDAR, columns={'ca_dt': ca_dt})
+      FILTER(condition=pr_name == 'RubyCopper-Star':string, columns={'pr_release': pr_release})
+       SCAN(table=main.PRODUCTS, columns={'pr_name': pr_name, 'pr_release': pr_release})
+     SCAN(table=main.DEVICES, columns={'de_product_id': de_product_id, 'de_purchase_ts': de_purchase_ts})
+    FILTER(condition=pr_name == 'RubyCopper-Star':string, columns={'pr_id': pr_id})
+     SCAN(table=main.PRODUCTS, columns={'pr_id': pr_id, 'pr_name': pr_name})

tests/test_plan_refsols/lines_shipping_vs_customer_region.txt

@@ -1,5 +1,5 @@
-ROOT(columns=[('order_year', YEAR(o_orderdate)), ('customer_region_name', r_name), ('customer_nation_name', name_3), ('supplier_region_name', name_16), ('nation_name', n_name)], orderings=[])
- JOIN(condition=t0.l_partkey == t1.ps_partkey & t0.l_suppkey == t1.ps_suppkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_name': t1.n_name, 'name_16': t1.r_name, 'name_3': t0.n_name, 'o_orderdate': t0.o_orderdate, 'r_name': t0.r_name})
+ROOT(columns=[('order_year', YEAR(o_orderdate)), ('customer_region_name', r_name), ('customer_nation_name', nname), ('supplier_region_name', name_16), ('nation_name', n_name)], orderings=[])
+ JOIN(condition=t0.l_partkey == t1.ps_partkey & t0.l_suppkey == t1.ps_suppkey, type=INNER, cardinality=SINGULAR_ACCESS, columns={'n_name': t1.n_name, 'name_16': t1.r_name, 'nname': t0.n_name, 'o_orderdate': t0.o_orderdate, 'r_name': t0.r_name})
   JOIN(condition=t0.o_orderkey == t1.l_orderkey, type=INNER, cardinality=PLURAL_ACCESS, columns={'l_partkey': t1.l_partkey, 'l_suppkey': t1.l_suppkey, 'n_name': t0.n_name, 'o_orderdate': t0.o_orderdate, 'r_name': t0.r_name})
    JOIN(condition=t0.c_custkey == t1.o_custkey, type=INNER, cardinality=PLURAL_FILTER, columns={'n_name': t0.n_name, 'o_orderdate': t1.o_orderdate, 'o_orderkey': t1.o_orderkey, 'r_name': t0.r_name})
     JOIN(condition=t0.n_nationkey == t1.c_nationkey, type=INNER, cardinality=PLURAL_ACCESS, columns={'c_custkey': t1.c_custkey, 'n_name': t0.n_name, 'r_name': t0.r_name})