Fixing CROSS bug with partition edge cases (#423)

Fixing recently exposed bugs with how `CROSS` handles certain edge cases involving partitions. Resolutions for this bug:
- Adjusted the `all_terms` clause of a GlobalContext to include ancestrally inherited terms
- Adjusted qualification to remove the `is_cross` term and instead qualify the RHS of the cross with an _augmented_ version of the LHS qualified answer as the context, with a new GlobalContext appended to it
- Adjusted the relational conversion of the children of a partition clause to avoid deleting join keys that are NOT original grouping keys

Also added some new tests to trigger patterns related to this bug or the fixes.

Author: knassre-bodo

pydough/conversion/relational_converter.py

@@ -669,11 +669,20 @@ def handle_children(
                 join_keys: list[tuple[HybridExpr, HybridExpr]] | None = (
                     child.subtree.join_keys
                 )
+
+                # If handling the child of a partition, remove any join keys
+                # where the LHS is a simple reference, since those keys are
+                # guaranteed to be present due to the partitioning.
                 if (
                     isinstance(hybrid.pipeline[pipeline_idx], HybridPartition)
                     and child_idx == 0
                 ):
-                    join_keys = None
+                    new_join_keys: list[tuple[HybridExpr, HybridExpr]] = []
+                    if join_keys is not None:
+                        for lhs_key, rhs_key in join_keys:
+                            if not isinstance(lhs_key, HybridRefExpr):
+                                new_join_keys.append((lhs_key, rhs_key))
+                    join_keys = new_join_keys if len(new_join_keys) > 0 else None
                 child_expr: HybridExpr
                 match child.connection_type:
                     case (
@@ -1461,11 +1470,7 @@ def optimize_relational_tree(
     # Run projection pullup to move projections as far up the tree as possible.
     # This is done as soon as possible to make joins redundant if they only
     # exist to compute a scalar projection and then link it with the data.
-    # print()
-    # print(root.to_tree_string())
     root = confirm_root(pullup_projections(root))
-    # print()
-    # print(root.to_tree_string())
 
     # Push filters down as far as possible
     root = confirm_root(push_filters(root, configs))

pydough/qdag/collections/global_context.py

@@ -13,7 +13,7 @@
 )
 from pydough.qdag.abstract_pydough_qdag import PyDoughQDAG
 from pydough.qdag.errors import PyDoughQDAGException
-from pydough.qdag.expressions import CollationExpression
+from pydough.qdag.expressions import BackReferenceExpression, CollationExpression
 
 from .collection_qdag import PyDoughCollectionQDAG
 from .collection_tree_form import CollectionTreeForm
@@ -33,6 +33,7 @@ def __init__(
         self._graph = graph
         self._collections: dict[str, PyDoughCollectionQDAG] = {}
         self._ancestral_mapping: dict[str, int] = {}
+        self._all_terms: set[str] = set()
         # If this collection has an ancestor, inherit its ancestors
         # and update depths, to preserve sub-collection context.
         # This ensures that downstream operations have the correct hierarchy.
@@ -44,11 +45,13 @@ def __init__(
                         f"Name {name!r} conflicts with a collection name in the graph {graph.name!r}"
                     )
                 else:
+                    self._all_terms.add(name)
                     self._ancestral_mapping[name] = level + 1
         for collection_name in graph.get_collection_names():
             meta = graph.get_collection(collection_name)
             assert isinstance(meta, CollectionMetadata)
             self._collections[collection_name] = TableCollection(meta, self)
+            self._all_terms.add(collection_name)
 
     @property
     def graph(self) -> GraphMetadata:
@@ -98,7 +101,7 @@ def inherited_downstreamed_terms(self) -> set[str]:
 
     @property
     def all_terms(self) -> set[str]:
-        return set(self.collections)
+        return self._all_terms
 
     @property
     def ordering(self) -> list[CollationExpression] | None:
@@ -115,11 +118,22 @@ def get_expression_position(self, expr_name: str) -> int:
         raise PyDoughQDAGException(f"Cannot call get_expression_position on {self!r}")
 
     def get_term(self, term_name: str) -> PyDoughQDAG:
-        if term_name not in self.collections:
+        if term_name in self.ancestral_mapping:
+            # Verify that the ancestor name is not also a name in the current
+            # context.
+            if term_name in self.collections:
+                raise PyDoughQDAGException(
+                    f"Cannot have term name {term_name!r} used in an ancestor of collection {self!r}"
+                )
+            # Create a back-reference to the ancestor term.
+            return BackReferenceExpression(
+                self, term_name, self.ancestral_mapping[term_name]
+            )
+        elif term_name in self.collections:
+            return self.collections[term_name]
+        else:
             raise PyDoughQDAGException(self.name_mismatch_error(term_name))
 
-        return self.collections[term_name]
-
     @property
     def standalone_string(self) -> str:
         return self.graph.name
@@ -157,4 +171,8 @@ def to_tree_form(self, is_last: bool) -> CollectionTreeForm:
             return self.to_tree_form_isolated(is_last)
 
     def equals(self, other: object) -> bool:
-        return isinstance(other, GlobalContext) and self.graph == other.graph
+        return (
+            isinstance(other, GlobalContext)
+            and self.graph == other.graph
+            and self.ancestor_context == other.ancestor_context
+        )