[BSE-4155] Add support for the Aggregate relational node

pydough/relational/__init__.py

@@ -3,10 +3,8 @@
 """
 
 __all__ = [
-    "Column",
     "Relational",
 ]
 from .abstract import (
-    Column,
     Relational,
 )

pydough/relational/abstract.py

@@ -1,17 +1,17 @@
 """
-This module contains the abstract base classes for the relational
+This file contains the abstract base classes for the relational
 representation. This roughly maps to a Relational Algebra representation
 but is not exact because it needs to maintain PyDough traits that define
 ordering and other properties of the relational expression.
 """
 
 from abc import ABC, abstractmethod
 from collections.abc import MutableMapping, MutableSequence
-from typing import Any, NamedTuple
+from typing import Any
 
-from sqlglot.expressions import Expression
+from sqlglot.expressions import Expression as SQLGlotExpression
 
-from pydough.pydough_ast.expressions import PyDoughExpressionAST
+from .relational_expressions import RelationalExpression
 
 
 class Relational(ABC):
@@ -20,6 +20,9 @@ class Relational(ABC):
     structure of all relational nodes in the PyDough system.
     """
 
+    def __init__(self, columns: MutableMapping[str, RelationalExpression]) -> None:
+        self._columns: MutableMapping[str, RelationalExpression] = columns
+
     @property
     @abstractmethod
     def inputs(self) -> MutableSequence["Relational"]:
@@ -32,55 +35,62 @@ def inputs(self) -> MutableSequence["Relational"]:
         """
 
     @property
-    def traits(self) -> MutableMapping[str, Any]:
+    def columns(self) -> MutableMapping[str, RelationalExpression]:
         """
-        Return the traits of the relational expression.
-        The traits in general may have a variable schema,
-        but each entry should be strongly defined. Here are
-        traits that should always be available:
+        Returns the columns of the relational expression.
 
-        - orderings: MutableSequence[PyDoughExpressionAST]
+        TODO: Associate an ordering in the future to avoid unnecessary SQL with the
+        final ordering of the root nodes.
 
         Returns:
-            MutableMapping[str, Any]: The traits of the relational expression.
+            MutableMapping[str, RelationalExpression]: The columns of the relational expression.
+                This does not have a defined ordering.
         """
-        return {"orderings": self.orderings}
+        return self._columns
 
-    @property
     @abstractmethod
-    def orderings(self) -> MutableSequence["PyDoughExpressionAST"]:
+    def node_equals(self, other: "Relational") -> bool:
         """
-        Returns the PyDoughExpressionAST that the relational expression is ordered by.
-        Each PyDoughExpressionAST is a result computed relative to the given set of columns.
+        Determine if two relational nodes are exactly identical,
+        excluding column ordering. This should be extended to avoid
+        duplicating equality logic shared across relational nodes.
+
+        Args:
+            other (Relational): The other relational node to compare against.
 
         Returns:
-            MutableSequence[PyDoughExpressionAST]: The PyDoughExpressionAST that the relational expression is ordered by,
-            possibly empty.
+            bool: Are the two relational nodes equal.
         """
 
-    @property
-    @abstractmethod
-    def columns(self) -> MutableSequence["Column"]:
+    def equals(self, other: "Relational") -> bool:
         """
-        Returns the columns of the relational expression.
+        Determine if two relational nodes are exactly identical,
+        including column ordering.
+
+        Args:
+            other (Relational): The other relational node to compare against.
 
         Returns:
-            MutableSequence[Column]: The columns of the relational expression.
+            bool: Are the two relational nodes equal.
         """
+        return self.node_equals(other) and self.columns == other.columns
+
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, Relational) and self.equals(other)
 
     @abstractmethod
-    def to_sqlglot(self) -> "Expression":
-        """Translate the given relational expression
+    def to_sqlglot(self) -> SQLGlotExpression:
+        """Translate the given relational node
         and its children to a SQLGlot expression.
 
         Returns:
-            Expression: A SqlGlot expression representing the relational expression.
+            Expression: A SqlGlot expression representing the relational node.
         """
 
     @abstractmethod
     def to_string(self) -> str:
         """
-        Convert the relational expression to a string.
+        Convert the relational node to a string.
 
         TODO: Refactor this API to include some form of string
         builder so we can draw lines between children properly.
@@ -90,38 +100,5 @@ def to_string(self) -> str:
             with this node at the root.
         """
 
-    @abstractmethod
-    def can_merge(self, other: "Relational") -> bool:
-        """
-        Determine if two relational nodes can be merged together.
-
-        Args:
-            other (Relational): The other relational node to merge with.
-
-        Returns:
-            bool: Can the two relational nodes be merged together.
-        """
-
-    @abstractmethod
-    def merge(self, other: "Relational") -> "Relational":
-        """
-        Merge two relational nodes together to produce one output
-        relational node. This requires can_merge to return True.
-
-        Args:
-            other (Relational): The other relational node to merge with.
-
-        Returns:
-            Relational: A new relational node that is the result of merging
-            the two input relational nodes together and removing any redundant
-            components.
-        """
-
-
-class Column(NamedTuple):
-    """
-    An column expression consisting of a name and an expression.
-    """
-
-    name: str
-    expr: "PyDoughExpressionAST"
+    def __repr__(self) -> str:
+        return self.to_string()

pydough/relational/aggregate.py

@@ -0,0 +1,64 @@
+"""
+This file contains the relational implementation for an aggregation. This is our
+relational representation for any grouping operation that optionally involves
+keys and aggregate functions.
+"""
+
+from collections.abc import MutableMapping
+
+from sqlglot.expressions import Expression as SQLGlotExpression
+
+from .abstract import Relational
+from .relational_expressions.call_expression import CallExpression
+from .relational_expressions.column_reference import ColumnReference
+from .single_relational import SingleRelational
+
+
+class Aggregate(SingleRelational):
+    """
+    The Aggregate node in the relational tree. This node represents an aggregation
+    based on some keys, which should most commonly be column references, and some
+    aggregate functions.
+    """
+
+    def __init__(
+        self,
+        input: Relational,
+        keys: MutableMapping[str, ColumnReference],
+        aggregations: MutableMapping[str, CallExpression],
+    ) -> None:
+        total_cols = {**keys, **aggregations}
+        assert len(total_cols) == len(keys) + len(
+            aggregations
+        ), "Keys and aggregations must have unique names"
+        super().__init__(input, total_cols)
+        self._keys: MutableMapping[str, ColumnReference] = keys
+        self._aggregations: MutableMapping[str, CallExpression] = aggregations
+        assert all(
+            agg.is_aggregation for agg in aggregations.values()
+        ), "All functions used in aggregations must be aggregation functions"
+
+    @property
+    def keys(self) -> MutableMapping[str, ColumnReference]:
+        return self._keys
+
+    @property
+    def aggregations(self) -> MutableMapping[str, CallExpression]:
+        return self._aggregations
+
+    def to_sqlglot(self) -> SQLGlotExpression:
+        raise NotImplementedError(
+            "Conversion to SQLGlot Expressions is not yet implemented."
+        )
+
+    def node_equals(self, other: Relational) -> bool:
+        return (
+            isinstance(other, Aggregate)
+            and self.keys == other.keys
+            and self.aggregations == other.aggregations
+            and super().node_equals(other)
+        )
+
+    def to_string(self) -> str:
+        # TODO: Should we visit the input?
+        return f"AGGREGATE(keys={self.keys}, aggregations={self.aggregations})"

pydough/relational/limit.py

@@ -0,0 +1,67 @@
+"""
+This file contains the relational implementation for a "limit" operation.
+This is the relational representation of top-n selection and typically depends
+on explicit ordering of the input relation.
+"""
+
+from collections.abc import MutableMapping, MutableSequence
+
+from sqlglot.expressions import Expression as SQLGlotExpression
+
+from pydough.types.integer_types import IntegerType
+
+from .abstract import Relational
+from .relational_expressions import ColumnOrdering, RelationalExpression
+from .single_relational import SingleRelational
+
+
+class Limit(SingleRelational):
+    """
+    The Limit node in the relational tree. This node represents any TOP-N
+    operations in the relational algebra. This operation is dependent on the
+    orderings of the input relation.
+    """
+
+    def __init__(
+        self,
+        input: Relational,
+        limit: RelationalExpression,
+        columns: MutableMapping[str, RelationalExpression],
+        orderings: MutableSequence[ColumnOrdering] | None = None,
+    ) -> None:
+        super().__init__(input, columns)
+        # Note: The limit is a relational expression because it should be a constant
+        # now but in the future could be a more complex expression that may require
+        # multi-step SQL to successfully evaluate.
+        assert isinstance(
+            limit.data_type, IntegerType
+        ), "Limit must be an integer type."
+        self._limit: RelationalExpression = limit
+        self._orderings: MutableSequence[ColumnOrdering] = (
+            [] if orderings is None else orderings
+        )
+
+    @property
+    def limit(self) -> RelationalExpression:
+        return self._limit
+
+    @property
+    def orderings(self) -> MutableSequence[ColumnOrdering]:
+        return self._orderings
+
+    def to_sqlglot(self) -> SQLGlotExpression:
+        raise NotImplementedError(
+            "Conversion to SQLGlot Expressions is not yet implemented."
+        )
+
+    def node_equals(self, other: Relational) -> bool:
+        return (
+            isinstance(other, Limit)
+            and self.limit == other.limit
+            and self.orderings == other.orderings
+            and super().node_equals(other)
+        )
+
+    def to_string(self) -> str:
+        # TODO: Should we visit the input?
+        return f"LIMIT(limit={self.limit}, columns={self.columns}, orderings={self.orderings})"

pydough/relational/project.py

@@ -0,0 +1,42 @@
+"""
+This file contains the relational implementation for a "project". This is our
+relational representation for a "calc" that involves any compute steps and can include
+adding or removing columns (as well as technically reordering). In general, we seek to
+avoid introducing extra nodes just to reorder or prune columns, so ideally their use
+should be sparse.
+"""
+
+from collections.abc import MutableMapping
+
+from sqlglot.expressions import Expression as SQLGlotExpression
+
+from .abstract import Relational
+from .relational_expressions import RelationalExpression
+from .single_relational import SingleRelational
+
+
+class Project(SingleRelational):
+    """
+    The Project node in the relational tree. This node represents a "calc" in
+    relational algebra, which should involve some "compute" functions and may
+    involve adding, removing, or reordering columns.
+    """
+
+    def __init__(
+        self,
+        input: Relational,
+        columns: MutableMapping[str, RelationalExpression],
+    ) -> None:
+        super().__init__(input, columns)
+
+    def to_sqlglot(self) -> SQLGlotExpression:
+        raise NotImplementedError(
+            "Conversion to SQLGlot Expressions is not yet implemented."
+        )
+
+    def node_equals(self, other: Relational) -> bool:
+        return isinstance(other, Project) and super().node_equals(other)
+
+    def to_string(self) -> str:
+        # TODO: Should we visit the input?
+        return f"PROJECT(columns={self.columns})"

pydough/relational/relational_expressions/__init__.py

@@ -0,0 +1,10 @@
+"""
+TODO: add module-level docstring
+"""
+
+__all__ = [
+    "ColumnOrdering",
+    "RelationalExpression",
+]
+from .abstract import RelationalExpression
+from .column_ordering import ColumnOrdering