add collection back

Author: hadia206

pydough/conversion/hybrid_translator.py

@@ -1118,26 +1118,22 @@ def make_hybrid_expr(
                 return HybridBackRefExpr(expr_name, expr.back_levels, expr.pydough_type)
             case Reference():
                 if hybrid.ancestral_mapping.get(expr.term_name, 0) > 0:
-                    # HA: FIXME: raise error for now.
-                    raise NotImplementedError(
-                        "Unable to resolve collection context for Reference. "
+                    collection = expr.collection
+                    while (
+                        isinstance(collection, PartitionChild)
+                        and expr.term_name in collection.child_access.ancestral_mapping
+                    ):
+                        collection = collection.child_access
+                    return self.make_hybrid_expr(
+                        hybrid,
+                        BackReferenceExpression(
+                            collection,
+                            expr.term_name,
+                            hybrid.ancestral_mapping[expr.term_name],
+                        ),
+                        child_ref_mapping,
+                        inside_agg,
                     )
-                    # collection = expr.collection
-                    # while (
-                    #     isinstance(collection, PartitionChild)
-                    #     and expr.term_name in collection.child_access.ancestral_mapping
-                    # ):
-                    #     collection = collection.child_access
-                    # return self.make_hybrid_expr(
-                    #     hybrid,
-                    #     BackReferenceExpression(
-                    #         collection,
-                    #         expr.term_name,
-                    #         hybrid.ancestral_mapping[expr.term_name],
-                    #     ),
-                    #     child_ref_mapping,
-                    #     inside_agg,
-                    # )
                 expr_name = hybrid.pipeline[-1].renamings.get(
                     expr.term_name, expr.term_name
                 )

pydough/conversion/relational_converter.py

@@ -1320,12 +1320,12 @@ def preprocess_root(
         if output_cols is None:
             for name in node.calc_terms:
                 name_typ: PyDoughType = node.get_expr(name).pydough_type
-                final_terms.append((name, Reference(name, name_typ)))
+                final_terms.append((name, Reference(node, name, name_typ)))
             final_terms.sort(key=lambda term: node.get_expression_position(term[0]))
         else:
             for _, column in output_cols:
                 column_typ: PyDoughType = node.get_expr(column).pydough_type
-                final_terms.append((column, Reference(column, column_typ)))
+                final_terms.append((column, Reference(node, column, column_typ)))
         children: list[PyDoughCollectionQDAG] = []
         final_calc: Calculate = Calculate(node, children).with_terms(final_terms)
         return final_calc

pydough/qdag/README.md

@@ -67,7 +67,7 @@ table_collection = builder.build_child_access("Nations", global_context_node)
 # Equivalent PyDough code: `TPCH.Nations.name`
 ref_name = "name"
 pydough_type = table_collection.get_expr(ref_name).pydough_type
-reference_node = builder.build_reference(ref_name, pydough_type)
+reference_node = builder.build_reference(table_collection, ref_name, pydough_type)
 
 # Build an expression function call node
 # Equivalent PyDough code: `LOWER(TPCH.Nations.name)`
@@ -104,7 +104,7 @@ nations_sub_collection = builder.build_child_access("nations", regions_collectio
 
 ref_name = "key"
 pydough_type = nations_sub_collection.get_expr(ref_name).pydough_type
-key_ref = builder.build_reference(ref_name, pydough_type)
+key_ref = builder.build_reference(nations_sub_collection, ref_name, pydough_type)
 literal_4 = builder.build_literal(4, NumericType())
 condition = builder.build_expression_function_call("EQU", [key_ref, literal_4])
 # Build WHERE node with condition
@@ -115,7 +115,7 @@ singular_node = builder.build_singular(where_node)
 # Build reference node for name
 ref_name = "name"
 pydough_type = singular_node.get_expr(ref_name).pydough_type
-reference_node = builder.build_reference(ref_name, pydough_type)
+reference_node = builder.build_reference(singular_node, ref_name, pydough_type)
 # Build CALCULATE node with calculated term
 calculate_node = builder.build_calc(regions_collection, [nations_sub_collection])
 calculate_node = calculate_node.with_terms([("n_4_nation", reference_node)])
@@ -139,7 +139,7 @@ top_k_node = top_k_node.with_collation([collation_expression])
 part_collection = builder.build_child_access("Parts", global_context_node)
 ref_name = "part_type"
 pydough_type = part_collection.get_expr(ref_name).pydough_type
-partition_key = builder.build_reference(ref_name, pydough_type)
+partition_key = builder.build_reference(part_collection, ref_name, pydough_type)
 partition_by_node = builder.build_partition(part_collection, child_collection, "p")
 partition_by_node = partition_by_node.with_keys([partition_key])
 

pydough/qdag/collections/augmenting_child_operator.py

@@ -85,7 +85,7 @@ def get_term(self, term_name: str) -> PyDoughQDAG:
             term = term.clone_with_parent(self)
         elif isinstance(term, PyDoughExpressionQDAG):
             typ = self.preceding_context.get_expr(term_name).pydough_type
-            term = Reference(term_name, typ)
+            term = Reference(self.preceding_context, term_name, typ)
         return term
 
     @cache

pydough/qdag/collections/collection_access.py

@@ -129,7 +129,9 @@ def get_term(self, term_name: str) -> PyDoughQDAG:
                 else:
                     assert context.ancestor_context is not None
                     context = context.ancestor_context
-            return Reference(term_name, context.get_expr(term_name).pydough_type)
+            return Reference(
+                context, term_name, context.get_expr(term_name).pydough_type
+            )
 
         if term_name not in self.all_terms:
             raise PyDoughQDAGException(self.name_mismatch_error(term_name))

pydough/qdag/collections/partition_child.py

@@ -102,7 +102,9 @@ def get_term(self, term_name: str):
                 else:
                     assert context.ancestor_context is not None
                     context = context.ancestor_context
-            return Reference(term_name, context.get_expr(term_name).pydough_type)
+            return Reference(
+                context, term_name, context.get_expr(term_name).pydough_type
+            )
 
         elif term_name not in self.all_terms:
             raise PyDoughQDAGException(self.name_mismatch_error(term_name))

pydough/qdag/expressions/back_reference_expression.py

@@ -40,13 +40,6 @@ def __init__(
         self._expression = self._ancestor.get_expr(term_name)
         self._term_type = self._expression.pydough_type
 
-    @property
-    def collection(self) -> PyDoughCollectionQDAG:
-        """
-        The collection that the Reference term comes from.
-        """
-        return self._collection
-
     @property
     def expression(self) -> PyDoughExpressionQDAG:
         """

pydough/qdag/expressions/child_reference_expression.py

@@ -35,13 +35,6 @@ def __init__(
                 f"Cannot reference plural expression {self.expression} from {self.collection}"
             )
 
-    @property
-    def collection(self) -> PyDoughCollectionQDAG:
-        """
-        The collection that the Reference term comes from.
-        """
-        return self._collection
-
     @property
     def expression(self) -> PyDoughExpressionQDAG:
         """

pydough/qdag/expressions/reference.py

@@ -7,6 +7,7 @@
 
 
 from pydough.qdag.abstract_pydough_qdag import PyDoughQDAG
+from pydough.qdag.collections.collection_qdag import PyDoughCollectionQDAG
 from pydough.types import PyDoughType
 
 from .expression_qdag import PyDoughExpressionQDAG
@@ -18,10 +19,20 @@ class Reference(PyDoughExpressionQDAG):
     a preceding collection.
     """
 
-    def __init__(self, term_name: str, term_type: PyDoughType):
+    def __init__(
+        self, collection: PyDoughCollectionQDAG, term_name: str, term_type: PyDoughType
+    ):
+        self._collection: PyDoughCollectionQDAG = collection
         self._term_name: str = term_name
         self._term_type: PyDoughType = term_type
 
+    @property
+    def collection(self) -> PyDoughCollectionQDAG:
+        """
+        The collection that the Reference term comes from.
+        """
+        return self._collection
+
     @property
     def term_name(self) -> str:
         """
@@ -51,5 +62,6 @@ def equals(self, other: object) -> bool:
         return (
             isinstance(other, Reference)
             and self.term_name == other.term_name
+            and self.collection.equals(other.collection)
             and self.pydough_type == other.pydough_type
         )

pydough/qdag/node_builder.py

@@ -155,11 +155,14 @@ def build_window_call(
             window_operator, qualified_args, collation_args, levels, kwargs
         )
 
-    def build_reference(self, name: str, typ: PyDoughType) -> Reference:
+    def build_reference(
+        self, collection: PyDoughCollectionQDAG, name: str, typ: PyDoughType
+    ) -> Reference:
         """
         Creates a new reference to an expression in the collection.
 
         Args:
+            `collection`: the collection that the reference comes from.
             `name`: the name of the expression being referenced.
             `typ`: the PyDough type of the expression being referenced.
 
@@ -170,7 +173,7 @@ def build_reference(self, name: str, typ: PyDoughType) -> Reference:
             `PyDoughQDAGException`: if `name` does not refer to an expression in
             the collection.
         """
-        return Reference(name, typ)
+        return Reference(collection, name, typ)
 
     def build_child_reference_expression(
         self,