range collections base, initial try

Author: hadia206

pydough/__init__.py

@@ -12,6 +12,7 @@
     "get_logger",
     "init_pydough_context",
     "parse_json_metadata_from_file",
+    "range_collection",
     "to_df",
     "to_sql",
 ]
@@ -22,6 +23,7 @@
 from .logger import get_logger
 from .metadata import parse_json_metadata_from_file
 from .unqualified import display_raw, from_string, init_pydough_context
+from .qdag.collections import range_collection
 
 # Create a default session for the user to interact with.
 # In most situations users will just use this session and

pydough/qdag/collections/__init__.py

@@ -21,6 +21,7 @@
     "TableCollection",
     "TopK",
     "Where",
+    "range_collection",
 ]
 
 from .augmenting_child_operator import AugmentingChildOperator

pydough/qdag/collections/range_collection.py

@@ -0,0 +1,44 @@
+"""A user-defined collection of integers in a specified range.
+Usage:
+`pydough.range_collection(start, end, step, name, column_name)`
+
+This module defines a collection that generates integers from `start` to `end`
+with a specified `step`. The user must specify the name of the collection and the
+name of the column that will hold the integer values.
+"""
+
+from pydough.qdag import PyDoughQDAG
+from pydough.types import NumericType
+from pydough.types.pydough_type import PyDoughType
+
+from .user_collections import PyDoughUserGeneratedCollection
+
+all = ["RangeGeneratedCollection"]
+
+
+class RangeGeneratedCollection(PyDoughUserGeneratedCollection):
+    """Integer range-based collection."""
+
+    def __init__(
+        self,
+        name: str,
+        column_name: str,
+        start: PyDoughQDAG,
+        end: PyDoughQDAG,
+        step: PyDoughQDAG,
+    ) -> None:
+        super().__init__(name=name, columns=[column_name])
+        self.start = start
+        self.end = end
+        self.step = step
+
+    @property
+    def column_names_and_types(self) -> list[tuple[str, PyDoughType]]:
+        return [(self.columns[0], NumericType())]
+
+    def is_empty(self) -> bool:
+        pass  # return self.start < self.end
+        return False
+
+    def to_string(self) -> str:
+        return f"RangeCollection({self.name}: {self.columns[0]} from {self.start} to {self.end} step {self.step})"

pydough/qdag/collections/user_collection_qdag.py

@@ -0,0 +1,117 @@
+from pydough.qdag import PyDoughCollectionQDAG
+from pydough.qdag.abstract_pydough_qdag import PyDoughQDAG
+from pydough.qdag.errors import PyDoughQDAGException
+from pydough.qdag.expressions.collation_expression import CollationExpression
+from pydough.qdag.expressions.reference import Reference
+
+from .child_access import ChildAccess
+from .collection_tree_form import CollectionTreeForm
+from .user_collections import PyDoughUserGeneratedCollection
+
+
+# or should it be CollectionAccess?
+class PyDoughUserGeneratedCollectionQDag(ChildAccess):
+    def __init__(
+        self,
+        ancestor: PyDoughCollectionQDAG,
+        collection: PyDoughUserGeneratedCollection,
+    ):
+        assert ancestor is not None
+        super().__init__(ancestor)
+        self._collection = collection
+
+    def clone_with_parent(
+        self, new_ancestor: PyDoughCollectionQDAG
+    ) -> "PyDoughUserGeneratedCollectionQDag":
+        """
+        Copies `self` but with a new ancestor node that presumably has the
+        original ancestor in its predecessor chain.
+
+        Args:
+            `new_ancestor`: the node to use as the new parent of the clone.
+
+        Returns:
+            The cloned version of `self`.
+        """
+        return PyDoughUserGeneratedCollectionQDag(new_ancestor, self._collection)
+
+    def to_tree_form(self, is_last: bool) -> CollectionTreeForm:
+        raise NotImplementedError
+
+    @property
+    def collection(self) -> PyDoughUserGeneratedCollection:
+        """
+        The metadata for the table that is being referenced by the collection
+        node.
+        """
+        return self._collection
+
+    @property
+    def name(self) -> str:
+        return self.collection.name
+
+    @property
+    def preceding_context(self) -> PyDoughCollectionQDAG | None:
+        return None
+
+    @property
+    def ordering(self) -> list[CollationExpression] | None:
+        return None
+
+    @property
+    def calc_terms(self) -> set[str]:
+        return set(self.collection.columns)
+
+    @property
+    def all_terms(self) -> set[str]:
+        """
+        The set of expression/subcollection names accessible by the context.
+        """
+        return self.calc_terms
+
+    @property
+    def ancestral_mapping(self) -> dict[str, int]:
+        return self._ancestor.ancestral_mapping
+
+    @property
+    def inherited_downstreamed_terms(self) -> set[str]:
+        if self._ancestor:
+            return self._ancestor.inherited_downstreamed_terms
+        else:
+            return set()
+
+    def is_singular(self, context: "PyDoughCollectionQDAG") -> bool:
+        return False
+
+    def get_term(self, term_name: str) -> PyDoughQDAG:
+        if term_name not in self.collection.columns:
+            raise PyDoughQDAGException(self.name_mismatch_error(term_name))
+
+        return Reference(self._ancestor, term_name)
+
+    def get_expression_position(self, expr_name: str) -> int:
+        raise PyDoughQDAGException(f"Cannot call get_expression_position on {self!r}")
+
+    @property
+    def unique_terms(self) -> list[str]:
+        return self.collection.columns
+
+    @property
+    def standalone_string(self) -> str:
+        """
+        Returns a string representation of the collection in a standalone form.
+        This is used for debugging and logging purposes.
+        """
+        return f"UserGeneratedCollection({self.name}, {', '.join(self.collection.columns)})"
+
+    @property
+    def key(self) -> str:
+        return f"USER_GENERATED_COLLECTION-{self.name}"
+
+    def to_string(self) -> str:
+        # Stringify as "name(column_name)
+        return f"{self.name}({', '.join(self.collection.columns)})"
+
+    @property
+    def tree_item_string(self) -> str:
+        return f"UserGeneratedCollection({self.name})"

pydough/qdag/collections/user_collections.py

@@ -0,0 +1,61 @@
+"""
+Base definition of PyDough QDAG collection type for accesses to a user defined
+collection of the current context.
+"""
+
+from abc import ABC, abstractmethod
+
+from pydough.types.pydough_type import PyDoughType
+
+__all__ = ["PyDoughUserGeneratedCollection"]
+
+
+class PyDoughUserGeneratedCollection(ABC):
+    """
+    Abstract base class for a user defined table collection.
+    This class defines the interface for accessing a user defined table collection
+    directly, without any specific implementation details.
+    It is intended to be subclassed by specific implementations that provide
+    the actual behavior and properties of the collection.
+    """
+
+    def __init__(self, name: str, columns: list[str]) -> None:
+        self._name = name
+        self._columns = columns
+
+    def __eq__(self, other) -> bool:
+        return isinstance(other, PyDoughUserGeneratedCollection) and repr(self) == repr(
+            other
+        )
+
+    def __repr__(self) -> str:
+        return self.to_string()
+
+    def __hash__(self) -> int:
+        return hash(repr(self))
+
+    def __str__(self) -> str:
+        return self.to_string()
+
+    @property
+    def name(self) -> str:
+        """Return the name used for the collection."""
+        return self._name
+
+    @property
+    def columns(self) -> list[str]:
+        """Return column names."""
+        return self._columns
+
+    @property
+    @abstractmethod
+    def column_names_and_types(self) -> list[tuple[str, PyDoughType]]:
+        """Return column names and their types."""
+
+    @abstractmethod
+    def is_empty(self) -> bool:
+        """Check if the collection is empty."""
+
+    @abstractmethod
+    def to_string(self) -> str:
+        """Return a string representation of the collection."""

pydough/qdag/node_builder.py

@@ -18,6 +18,11 @@
     PyDoughOperator,
     builtin_registered_operators,
 )
+from pydough.qdag.collections.range_collection import RangeGeneratedCollection
+from pydough.qdag.collections.user_collection_qdag import (
+    PyDoughUserGeneratedCollectionQDag,
+)
+from pydough.qdag.collections.user_collections import PyDoughUserGeneratedCollection
 from pydough.types import PyDoughType
 
 from .abstract_pydough_qdag import PyDoughQDAG
@@ -393,3 +398,37 @@ def build_singular(
             The newly created PyDough SINGULAR instance.
         """
         return Singular(preceding_context)
+
+    def build_generated_collection(
+        self,
+        preceding_context: PyDoughCollectionQDAG,
+        name: str,
+        column_name: list[str],
+        args: list[PyDoughQDAG],
+    ) -> PyDoughUserGeneratedCollectionQDag:
+        """
+        Creates a new user-defined collection.
+
+        Args:
+            `name`: the name of the collection.
+            `column_name`: the name of the column that will hold the integer values.
+            `args`: the arguments that define the collection, such as a range of
+            integers or a dataframe.
+
+        Returns:
+            The newly created user-defined collection.
+        """
+        # TODO: case range vs. dataframe
+        if len(args) != 3:
+            raise PyDoughQDAGException(
+                f"Expected 3 arguments for range collection, got {len(args)}"
+            )
+        collection: PyDoughUserGeneratedCollection = RangeGeneratedCollection(
+            name, column_name[0], args[0], args[1], args[2]
+        )
+        collection_qdag: PyDoughUserGeneratedCollectionQDag = (
+            PyDoughUserGeneratedCollectionQDag(
+                ancestor=preceding_context, collection=collection
+            )
+        )
+        return collection_qdag

pydough/unqualified/qualification.py

@@ -38,6 +38,9 @@
     Where,
     WindowCall,
 )
+from pydough.qdag.collections.user_collection_qdag import (
+    PyDoughUserGeneratedCollectionQDag,
+)
 from pydough.types import PyDoughType
 
 from .errors import PyDoughUnqualifiedException
@@ -48,6 +51,7 @@
     UnqualifiedCalculate,
     UnqualifiedCollation,
     UnqualifiedCross,
+    UnqualifiedGeneratedCollection,
     UnqualifiedLiteral,
     UnqualifiedNode,
     UnqualifiedOperation,
@@ -1261,6 +1265,48 @@ def qualify_cross(
         )
         return qualified_child
 
+    def qualify_generated_collection(
+        self,
+        unqualified: UnqualifiedGeneratedCollection,
+        context: PyDoughCollectionQDAG,
+        is_child: bool,
+        is_cross: bool,
+    ) -> PyDoughUserGeneratedCollectionQDag:
+        """
+        Transforms an `UnqualifiedGeneratedCollection` into a PyDoughCollectionQDAG node.
+
+        Args:
+            `unqualified`: the UnqualifiedGeneratedCollection instance to be transformed.
+            `context`: the collection QDAG whose context the collection is being
+            evaluated within.
+            `is_child`: whether the collection is being qualified as a child
+            of a child operator context, such as CALCULATE or PARTITION.
+            `is_cross`: whether the collection being qualified is a CROSS JOIN operation
+
+        Returns:
+            The PyDough QDAG object for the qualified collection node.
+
+        """
+        qualified_args: list[PyDoughQDAG] = []
+
+        # TODO: how to handle dataframe case.
+        for arg in unqualified._parcel[2]:
+            assert isinstance(arg, UnqualifiedNode)
+            qualified_arg: PyDoughQDAG = self.qualify_node(
+                arg, context, [], is_child, is_cross
+            )
+            qualified_args.append(qualified_arg)
+
+        generated_collection_qdag: PyDoughUserGeneratedCollectionQDag = (
+            self.builder.build_generated_collection(
+                context,
+                unqualified._parcel[0],
+                unqualified._parcel[1],
+                qualified_args,
+            )
+        )
+        return generated_collection_qdag  # .collection
+
     def qualify_node(
         self,
         unqualified: UnqualifiedNode,
@@ -1333,6 +1379,10 @@ def qualify_node(
                 answer = self.qualify_best(unqualified, context, is_child, is_cross)
             case UnqualifiedCross():
                 answer = self.qualify_cross(unqualified, context, is_child, is_cross)
+            case UnqualifiedGeneratedCollection():
+                answer = self.qualify_generated_collection(
+                    unqualified, context, is_child, is_cross
+                )
             case _:
                 raise PyDoughUnqualifiedException(
                     f"Cannot qualify {unqualified.__class__.__name__}: {unqualified!r}"