Add support for uses of BACK that cause correlated references

pydough/conversion/relational_converter.py

@@ -27,6 +27,7 @@
     CallExpression,
     ColumnPruner,
     ColumnReference,
+    CorrelatedReference,
     EmptySingleton,
     ExpressionSortInfo,
     Filter,
@@ -43,6 +44,7 @@
 )
 from pydough.types import BooleanType, Int64Type, UnknownType
 
+from .hybrid_decorrelater import run_hybrid_decorrelation
 from .hybrid_tree import (
     ConnectionType,
     HybridBackRefExpr,
@@ -52,6 +54,7 @@
     HybridCollectionAccess,
     HybridColumnExpr,
     HybridConnection,
+    HybridCorrelExpr,
     HybridExpr,
     HybridFilter,
     HybridFunctionExpr,
@@ -90,11 +93,20 @@ class TranslationOutput:
     value of that expression.
     """
 
+    correlated_name: str | None = None
+    """
+    The name that can be used to refer to the relational output in correlated
+    references.
+    """
+
 
 class RelTranslation:
     def __init__(self):
         # An index used for creating fake column names
         self.dummy_idx = 1
+        # A stack of contexts used to point to ancestors for correlated
+        # references.
+        self.stack: list[TranslationOutput] = []
 
     def make_null_column(self, relation: RelationalNode) -> ColumnReference:
         """
@@ -145,6 +157,24 @@ def get_column_name(
             new_name = f"{name}_{self.dummy_idx}"
         return new_name
 
+    def get_correlated_name(self, context: TranslationOutput) -> str:
+        """
+        Finds the name used to refer to a context for correlated variable
+        access. If the context does not have a correlated name, a new one is
+        generated for it.
+
+        Args:
+            `context`: the context containing the relational subtree being
+            referrenced in a correlated variable access.
+
+        Returns:
+            The name used to refer to the context in a correlated reference.
+        """
+        if context.correlated_name is None:
+            context.correlated_name = f"corr{self.dummy_idx}"
+            self.dummy_idx += 1
+        return context.correlated_name
+
     def translate_expression(
         self, expr: HybridExpr, context: TranslationOutput | None
     ) -> RelationalExpression:
@@ -199,8 +229,32 @@ def translate_expression(
                     order_inputs,
                     expr.kwargs,
                 )
+            case HybridCorrelExpr():
+                # Convert correlated expressions by converting the expression
+                # they point to in the context of the top of the stack, then
+                # wrapping the result in a correlated reference.
+                ancestor_context: TranslationOutput = self.stack.pop()
+                ancestor_expr: RelationalExpression = self.translate_expression(
+                    expr.expr, ancestor_context
+                )
+                self.stack.append(ancestor_context)
+                match ancestor_expr:
+                    case ColumnReference():
+                        return CorrelatedReference(
+                            ancestor_expr.name,
+                            self.get_correlated_name(ancestor_context),
+                            expr.typ,
+                        )
+                    case CorrelatedReference():
+                        return ancestor_expr
+                    case _:
+                        raise ValueError(
+                            f"Unsupported expression to reference in a correlated reference: {ancestor_expr}"
+                        )
             case _:
-                raise NotImplementedError(expr.__class__.__name__)
+                raise NotImplementedError(
+                    f"TODO: support relational conversion on {expr.__class__.__name__}"
+                )
 
     def join_outputs(
         self,
@@ -257,6 +311,7 @@ def join_outputs(
             [LiteralExpression(True, BooleanType())],
             [join_type],
             join_columns,
+            correl_name=lhs_result.correlated_name,
         )
         input_aliases: list[str | None] = out_rel.default_input_aliases
 
@@ -397,9 +452,11 @@ def handle_children(
         """
         for child_idx, child in enumerate(hybrid.children):
             if child.required_steps == pipeline_idx:
+                self.stack.append(context)
                 child_output = self.rel_translation(
                     child, child.subtree, len(child.subtree.pipeline) - 1
                 )
+                self.stack.pop()
                 assert child.subtree.join_keys is not None
                 join_keys: list[tuple[HybridExpr, HybridExpr]] = child.subtree.join_keys
                 agg_keys: list[HybridExpr]
@@ -592,7 +649,7 @@ def translate_partition(
 
         Returns:
             The TranslationOutput payload containing access to the aggregated
-            child corresponding tot he partition data.
+            child corresponding to the partition data.
         """
         expressions: dict[HybridExpr, ColumnReference] = {}
         # Account for the fact that the PARTITION is stepping down a level,
@@ -828,11 +885,26 @@ def rel_translation(
                 if isinstance(operation.collection, TableCollection):
                     result = self.build_simple_table_scan(operation)
                     if context is not None:
+                        # If the collection access is the child of something
+                        # else, join it onto that something else. Use the
+                        # uniqueness keys of the ancestor, which should also be
+                        # present in the collection (e.g. joining a partition
+                        # onto the original data using the partition keys).
+                        assert preceding_hybrid is not None
+                        join_keys: list[tuple[HybridExpr, HybridExpr]] = []
+                        for unique_column in sorted(
+                            preceding_hybrid[0].pipeline[0].unique_exprs, key=str
+                        ):
+                            if unique_column not in result.expressions:
+                                raise ValueError(
+                                    f"Cannot connect parent context to child {operation.collection} because {unique_column} is not in the child's expressions."
+                                )
+                            join_keys.append((unique_column, unique_column))
                         result = self.join_outputs(
                             context,
                             result,
                             JoinType.INNER,
-                            [],
+                            join_keys,
                             None,
                         )
                 else:
@@ -856,7 +928,8 @@ def rel_translation(
                 assert context is not None, "Malformed HybridTree pattern."
                 result = self.translate_filter(operation, context)
             case HybridPartition():
-                assert context is not None, "Malformed HybridTree pattern."
+                if context is None:
+                    context = TranslationOutput(EmptySingleton(), {})
                 result = self.translate_partition(
                     operation, context, hybrid, pipeline_idx
                 )
@@ -942,10 +1015,11 @@ def convert_ast_to_relational(
     final_terms: set[str] = node.calc_terms
     node = translator.preprocess_root(node)
 
-    # Convert the QDAG node to the hybrid form, then invoke the relational
-    # conversion procedure. The first element in the returned list is the
-    # final rel node.
-    hybrid: HybridTree = HybridTranslator(configs).make_hybrid_tree(node)
+    # Convert the QDAG node to the hybrid form, decorrelate it, then invoke
+    # the relational conversion procedure. The first element in the returned
+    # list is the final rel node.
+    hybrid: HybridTree = HybridTranslator(configs).make_hybrid_tree(node, None)
+    run_hybrid_decorrelation(hybrid)
     renamings: dict[str, str] = hybrid.pipeline[-1].renamings
     output: TranslationOutput = translator.rel_translation(
         None, hybrid, len(hybrid.pipeline) - 1

pydough/pydough_operators/base_operator.py

@@ -80,3 +80,15 @@ def to_string(self, arg_strings: list[str]) -> str:
         Returns:
             The string representation of the operator called on its arguments.
         """
+
+    @abstractmethod
+    def equals(self, other: object) -> bool:
+        """
+        Returns whether this operator is equal to another operator.
+        """
+
+    def __eq__(self, other: object) -> bool:
+        return self.equals(other)
+
+    def __hash__(self) -> int:
+        return hash(repr(self))

pydough/relational/README.md

@@ -18,6 +18,8 @@ The relational_expressions submodule provides functionality to define and manage
 - `ExpressionSortInfo`: The representation of ordering for an expression within a relational node.
 - `RelationalExpressionVisitor`: The basic Visitor pattern to perform operations across the expression components of a relational tree.
 - `ColumnReferenceFinder`: Finds all unique column references in a relational expression.
+- `CorrelatedReference`: The expression implementation for accessing a correlated column reference in a relational node.
+- `CorrelatedReferenceFinder`: Finds all unique correlated references in a relational expression.
 - `RelationalExpressionShuttle`: Specialized form of the visitor pattern that returns a relational expression.
 - `ColumnReferenceInputNameModifier`: Shuttle implementation designed to update all uses of a column reference's input name to a new input name based on a dictionary.
 
@@ -33,6 +35,7 @@ from pydough.relational.relational_expressions import (
     ExpressionSortInfo,
     ColumnReferenceFinder,
     ColumnReferenceInputNameModifier,
+    CorrelatedReferenceFinder,
     WindowCallExpression,
 )
 from pydough.pydough_operators import ADD, RANKING
@@ -64,6 +67,11 @@ unique_column_refs = finder.get_column_references()
 # Modify the input name of column references in the call expression
 modifier = ColumnReferenceInputNameModifier({"old_input_name": "new_input_name"})
 modified_call_expr = call_expr.accept_shuttle(modifier)
+
+# Find all unique correlated references in the call expression
+correlated_finder = CorrelatedReferenceFinder()
+call_expr.accept(correlated_finder)
+unique_correlated_refs = correlated_finder.get_correlated_references()
 ```
 
 ## [Relational Nodes](relational_nodes/README.md)

pydough/relational/__init__.py

@@ -6,6 +6,7 @@
     "ColumnReferenceFinder",
     "ColumnReferenceInputNameModifier",
     "ColumnReferenceInputNameRemover",
+    "CorrelatedReference",
     "EmptySingleton",
     "ExpressionSortInfo",
     "Filter",
@@ -30,6 +31,7 @@
     ColumnReferenceFinder,
     ColumnReferenceInputNameModifier,
     ColumnReferenceInputNameRemover,
+    CorrelatedReference,
     ExpressionSortInfo,
     LiteralExpression,
     RelationalExpression,

pydough/relational/relational_expressions/README.md

@@ -45,6 +45,14 @@ The relational_expressions module provides functionality to define and manage va
 
 - `ColumnReferenceFinder`: Finds all unique column references in a relational expression.
 
+### [correlated_reference.py](correlated_reference.py)
+
+- `CorrelatedReference`: The expression implementation for accessing a correlated column reference in a relational node.
+
+### [correlated_reference_finder.py](correlated_reference_finder.py)
+
+- `CorrelatedReferenceFinder`: Finds all unique correlated references in a relational expression.
+
 ### [relational_expression_shuttle.py](relational_expression_shuttle.py)
 
 - `RelationalExpressionShuttle`: Specialized form of the visitor pattern that returns a relational expression. This is used to handle the common case where we need to modify a type of input.
@@ -69,6 +77,8 @@ from pydough.relational.relational_expressions import (
     ExpressionSortInfo,
     ColumnReferenceFinder,
     ColumnReferenceInputNameModifier,
+    CorrelatedReference,
+    CorrelatedReferenceFinder,
 )
 from pydough.pydough_operators import ADD
 from pydough.types import Int64Type
@@ -82,6 +92,10 @@ literal_expr = LiteralExpression(10, Int64Type())
 # Create a call expression for addition
 call_expr = CallExpression(ADD, Int64Type(), [column_ref, literal_expr])
 
+# Create a correlated reference to column `column_name` in the first input to
+# an ancestor join of `corr1`
+correlated_ref = CorrelatedReference("column_name", "corr1", Int64Type())
+
 # Create an expression sort info
 sort_info = ExpressionSortInfo(call_expr, ascending=True, nulls_first=False)
 
@@ -96,4 +110,9 @@ unique_column_refs = finder.get_column_references()
 # Modify the input name of column references in the call expression
 modifier = ColumnReferenceInputNameModifier({"old_input_name": "new_input_name"})
 modified_call_expr = call_expr.accept_shuttle(modifier)
+
+# Find all unique correlated references in the call expression
+correlated_finder = CorrelatedReferenceFinder()
+call_expr.accept(correlated_finder)
+unique_correlated_refs = correlated_finder.get_correlated_references()
 ```

pydough/relational/relational_expressions/__init__.py

@@ -9,6 +9,8 @@
     "ColumnReferenceFinder",
     "ColumnReferenceInputNameModifier",
     "ColumnReferenceInputNameRemover",
+    "CorrelatedReference",
+    "CorrelatedReferenceFinder",
     "ExpressionSortInfo",
     "LiteralExpression",
     "RelationalExpression",
@@ -21,6 +23,8 @@
 from .column_reference_finder import ColumnReferenceFinder
 from .column_reference_input_name_modifier import ColumnReferenceInputNameModifier
 from .column_reference_input_name_remover import ColumnReferenceInputNameRemover
+from .correlated_reference import CorrelatedReference
+from .correlated_reference_finder import CorrelatedReferenceFinder
 from .expression_sort_info import ExpressionSortInfo
 from .literal_expression import LiteralExpression
 from .relational_expression_visitor import RelationalExpressionVisitor

pydough/relational/relational_expressions/abstract_expression.py

@@ -81,6 +81,9 @@ def to_string(self, compact: bool = False) -> str:
     def __repr__(self) -> str:
         return self.to_string()
 
+    def __hash__(self) -> int:
+        return hash(self.to_string())
+
     @abstractmethod
     def accept(self, visitor: RelationalExpressionVisitor) -> None:
         """

pydough/relational/relational_expressions/column_reference_finder.py

@@ -42,3 +42,6 @@ def visit_literal_expression(self, literal_expression: LiteralExpression) -> Non
 
     def visit_column_reference(self, column_reference: ColumnReference) -> None:
         self._column_references.add(column_reference)
+
+    def visit_correlated_reference(self, correlated_reference) -> None:
+        pass

pydough/relational/relational_expressions/column_reference_input_name_modifier.py

@@ -45,3 +45,6 @@ def visit_column_reference(self, column_reference) -> RelationalExpression:
             raise ValueError(
                 f"Input name {column_reference.input_name} not found in the input name map."
             )
+
+    def visit_correlated_reference(self, correlated_reference) -> RelationalExpression:
+        return correlated_reference

pydough/relational/relational_expressions/column_reference_input_name_remover.py

@@ -37,3 +37,6 @@ def visit_column_reference(self, column_reference) -> RelationalExpression:
                 column_reference.data_type,
                 None,
             )
+
+    def visit_correlated_reference(self, correlated_reference) -> RelationalExpression:
+        return correlated_reference