Set up PyDough collection nodes in AST module

pydough/pydough_ast/__init__.py

@@ -10,6 +10,10 @@
     "Literal",
     "ExpressionFunctionCall",
     "PyDoughASTException",
+    "PyDoughCollectionAST",
+    "TableCollection",
+    "SubCollection",
+    "Calc",
 ]
 
 from .abstract_pydough_ast import PyDoughAST
@@ -20,4 +24,5 @@
     Literal,
     ExpressionFunctionCall,
 )
+from .collections import PyDoughCollectionAST, TableCollection, Calc
 from .node_builder import AstNodeBuilder

pydough/pydough_ast/collections/README.md

@@ -0,0 +1 @@
+TODO: COMPLETE THIS README FOR THE PYDOUGH AST COLLECTIONS MODULE

pydough/pydough_ast/collections/__init__.py

@@ -0,0 +1,17 @@
+"""
+TODO: add module-level docstring
+"""
+
+__all__ = [
+    "PyDoughCollectionAST",
+    "TableCollection",
+    "Calc",
+    "GlobalContext",
+    "CollectionAccess",
+]
+
+from .collection_ast import PyDoughCollectionAST
+from .table_collection import TableCollection
+from .calc import Calc
+from .global_context import GlobalContext
+from .collection_access import CollectionAccess

pydough/pydough_ast/collections/calc.py

@@ -0,0 +1,126 @@
+"""
+TODO: add file-level docstring
+"""
+
+__all__ = ["Calc"]
+
+
+from typing import Dict, List, Tuple, Set
+
+from pydough.pydough_ast.abstract_pydough_ast import PyDoughAST
+from pydough.pydough_ast.errors import PyDoughASTException
+from pydough.pydough_ast.expressions import PyDoughExpressionAST
+from .collection_ast import PyDoughCollectionAST
+
+
+class Calc(PyDoughCollectionAST):
+    """
+    The AST node implementation class representing a CALC expression.
+    """
+
+    def __init__(
+        self,
+        predecessor: PyDoughCollectionAST,
+    ):
+        self._predecessor: PyDoughCollectionAST = predecessor
+        # Not initialized until with_terms is called
+        self._calc_term_indices: Dict[str, Tuple[int, PyDoughExpressionAST]] | None = (
+            None
+        )
+        self._all_terms: Dict[str, PyDoughExpressionAST] = {}
+
+    def with_terms(self, terms: List[Tuple[str, PyDoughExpressionAST]]) -> "Calc":
+        """
+        Specifies the terms that are calculated inside of a CALC node,
+        returning the mutated CALC node afterwards. This is called after the
+        CALC node is created so that the terms can be expressions that
+        reference child nodes of the CALC. However, this must be called
+        on the CALC node before any properties are accessed by `calc_terms`,
+        `all_terms`, `to_string`, etc.
+
+        Args:
+            `terms`: the list of terms calculated in the CALC node as a list of
+            tuples in the form `(name, expression)`. Each `expression` can
+            contain `ChildReference` instances that refer to an property of one
+            of the children of the CALC node.
+
+        Returns:
+            The mutated CALC node (which has also been modified in-place).
+
+        Raises:
+            `PyDoughASTException` if the terms have already been added to the
+            CALC node.
+        """
+        if self._calc_term_indices is not None:
+            raise PyDoughASTException(
+                "Cannot call `with_terms` on a CALC node more than once"
+            )
+        self._calc_term_indices = {name: idx for idx, (name, _) in enumerate(terms)}
+        # Include terms from the predecessor, with the terms from this CALC
+        # added in (overwriting any preceding properties with the same name)
+        self._all_terms: Dict[str, PyDoughExpressionAST] = {}
+        for name in self.preceding_context.all_terms:
+            self._all_terms[name] = self.preceding_context.get_term(name)
+        for name, property in terms:
+            self._all_terms[name] = property
+        return self
+
+    @property
+    def calc_term_indices(self) -> Dict[str, Tuple[int, PyDoughExpressionAST]]:
+        """
+        Mapping of each named expression of the CALC to a tuple (idx, expr)
+        where idx is the ordinal position of the property when included
+        in a CALC and property is the AST node representing the property.
+        """
+        if self._calc_term_indices is None:
+            raise PyDoughASTException(
+                "Cannot access `calc_term_indices` of a Calc node before adding calc terms with `with_terms`"
+            )
+        return self._calc_term_indices
+
+    @property
+    def ancestor_context(self) -> PyDoughCollectionAST | None:
+        return self._predecessor.ancestor_context
+
+    @property
+    def preceding_context(self) -> PyDoughCollectionAST | None:
+        return self._predecessor
+
+    @property
+    def calc_terms(self) -> Set[str]:
+        return set(self.calc_term_indices)
+
+    @property
+    def all_terms(self) -> Set[str]:
+        if self._calc_term_indices is None:
+            raise PyDoughASTException(
+                "Cannot access `all_terms` of a Calc node before adding calc terms with `with_terms`"
+            )
+        return set(self._all_terms)
+
+    def get_expression_position(self, expr_name: str) -> int:
+        if expr_name not in self.calc_term_indices:
+            raise PyDoughASTException(f"Unrecognized CALC term: {expr_name!r}")
+        return self.calc_term_indices[expr_name]
+
+    def get_term(self, term_name: str) -> PyDoughAST:
+        if term_name not in self.all_terms:
+            raise PyDoughASTException(f"Unrecognized term: {term_name!r}")
+        return self._all_terms[term_name]
+
+    def to_string(self) -> str:
+        kwarg_strings: List[str] = []
+        for name in self._calc_term_indices:
+            expr: PyDoughExpressionAST = self.get_term(name)
+            kwarg_strings.append(f"{name}={expr.to_string()}")
+        return f"{self.preceding_context.to_string()}({', '.join(kwarg_strings)})"
+
+    def to_tree_form(self) -> None:
+        raise NotImplementedError
+
+    def equals(self, other: "Calc") -> bool:
+        return (
+            super().equals(other)
+            and self.preceding_context == other.preceding_context
+            and self.terms == other.terms
+        )

pydough/pydough_ast/collections/collection_access.py

@@ -0,0 +1,119 @@
+"""
+TODO: add file-level docstring
+"""
+
+__all__ = ["CollectionAccess"]
+
+
+from typing import Dict, List, Tuple, Set
+
+from pydough.metadata import (
+    CollectionMetadata,
+    PropertyMetadata,
+)
+from pydough.pydough_ast.abstract_pydough_ast import PyDoughAST
+from pydough.pydough_ast.errors import PyDoughASTException
+from pydough.pydough_ast.expressions import ColumnProperty
+from .collection_ast import PyDoughCollectionAST
+
+
+class CollectionAccess(PyDoughCollectionAST):
+    """
+    The AST node implementation class representing a table collection accessed
+    either directly or as a subcollection of another collection.
+    """
+
+    def __init__(
+        self,
+        collection: CollectionMetadata,
+        ancestor: PyDoughCollectionAST,
+        predecessor: PyDoughCollectionAST | None = None,
+    ):
+        self._collection: CollectionMetadata = collection
+        self._ancestor: PyDoughCollectionAST = ancestor
+        self._predecessor: PyDoughCollectionAST | None = predecessor
+        self._properties: Dict[str, Tuple[int | None, PyDoughAST]] | None = None
+        self._calc_counter: int = 0
+
+    @property
+    def collection(self) -> CollectionMetadata:
+        """
+        The metadata for the table that is being referenced by the collection
+        node.
+        """
+        return self._collection
+
+    @property
+    def properties(self) -> Dict[str, Tuple[int | None, PyDoughAST]]:
+        """
+        Mapping of each property of the table to a tuple (idx, property)
+        where idx is the ordinal position of the property when included
+        in a CALC (None for subcollections), and property is the AST node
+        representing the property.
+
+        The properties are evaluated lazily & cached to prevent ping-ponging
+        between two tables that consider each other subcollections.
+        """
+
+        if self._properties is None:
+            self._properties = {}
+            # Ensure the properties are added in the same order they were
+            # defined in the metadata in to ensure dependencies are handled.
+            ordered_properties: List[str] = sorted(
+                self.collection.get_property_names(),
+                key=lambda p: self.collection.definition_order[p],
+            )
+            for property_name in ordered_properties:
+                property: PropertyMetadata = self.collection.get_property(property_name)
+                calc_idx: int | None
+                expression: PyDoughAST
+                if property.is_subcollection:
+                    # TODO: implement subcollections properly
+                    continue
+                else:
+                    calc_idx = self._calc_counter
+                    expression = ColumnProperty(property)
+                    self._calc_counter += 1
+                self._properties[property_name] = (calc_idx, expression)
+        return self._properties
+
+    @property
+    def ancestor_context(self) -> PyDoughCollectionAST | None:
+        return self._ancestor
+
+    @property
+    def preceding_context(self) -> PyDoughCollectionAST | None:
+        return self._predecessor
+
+    @property
+    def calc_terms(self) -> Set[str]:
+        # The calc terms are just all of the column properties (the ones
+        # that have an index)
+        return {name for name, (idx, _) in self.properties.items() if idx is not None}
+
+    @property
+    def all_terms(self) -> Set[str]:
+        return set(self.properties)
+
+    def get_expression_position(self, expr_name: str) -> int:
+        if expr_name not in self.properties:
+            raise PyDoughASTException(
+                f"Unrecognized term of {self.collection.error_name}: {expr_name!r}"
+            )
+        idx, _ = self.properties[expr_name]
+        if idx is None:
+            raise PyDoughASTException(
+                f"Cannot call get_expression_position on non-CALC term: {expr_name!r}"
+            )
+        return idx
+
+    def get_term(self, term_name: str) -> PyDoughAST:
+        if term_name not in self.properties:
+            raise PyDoughASTException(
+                f"Unrecognized term of {self.collection.error_name}: {term_name!r}"
+            )
+        _, term = self.properties[term_name]
+        return term
+
+    def equals(self, other: "CollectionAccess") -> bool:
+        return super().equals(other) and self.collection == other.collection