add range_collection to pydough top

Author: hadia206

pydough/qdag/collections/range_collection.py

@@ -1,13 +1,13 @@
 """A user-defined collection of integers in a specified range.
 Usage:
-`pydough.range_collection(start, end, step, name, column_name)`
+`pydough.range_collection(name, column, *args)`
+    args: start, end, step
 
 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
 
@@ -19,13 +19,14 @@
 class RangeGeneratedCollection(PyDoughUserGeneratedCollection):
     """Integer range-based collection."""
 
+    # HA_Q: should start/end/step be int or PyDoughQDAG? Why?
     def __init__(
         self,
         name: str,
         column_name: str,
-        start: PyDoughQDAG,
-        end: PyDoughQDAG,
-        step: PyDoughQDAG,
+        start: int,
+        end: int,
+        step: int,
     ) -> None:
         super().__init__(name=name, columns=[column_name])
         self.start = start
@@ -36,9 +37,14 @@ def __init__(
     def column_names_and_types(self) -> list[tuple[str, PyDoughType]]:
         return [(self.columns[0], NumericType())]
 
+    def __len__(self) -> int:
+        if self.start >= self.end:
+            return 0
+        return (self.end - self.start + self.step - 1) // self.step
+
     def is_empty(self) -> bool:
-        pass  # return self.start < self.end
-        return False
+        """Check if the range collection is empty."""
+        return len(self) == 0
 
     def to_string(self) -> str:
         return f"RangeCollection({self.name}: {self.columns[0]} from {self.start} to {self.end} step {self.step})"

pydough/qdag/node_builder.py

@@ -404,7 +404,7 @@ def build_generated_collection(
         preceding_context: PyDoughCollectionQDAG,
         name: str,
         column_name: list[str],
-        args: list[PyDoughQDAG],
+        args: list[int],
     ) -> PyDoughUserGeneratedCollectionQDag:
         """
         Creates a new user-defined collection.

pydough/unqualified/qualification.py

@@ -1287,22 +1287,22 @@ def qualify_generated_collection(
             The PyDough QDAG object for the qualified collection node.
 
         """
-        qualified_args: list[PyDoughQDAG] = []
+        # 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)
+        # # 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,
+                unqualified._parcel[2],
             )
         )
         return generated_collection_qdag  # .collection

pydough/unqualified/unqualified_node.py

@@ -450,11 +450,12 @@ def range_collection(
         """
         Method used to create a user-generated range collection node.
         """
-        range_args: list[UnqualifiedNode] = [
-            self.coerce_to_unqualified(start),
-            self.coerce_to_unqualified(stop),
-            self.coerce_to_unqualified(step),
-        ]
+        # range_args: list[UnqualifiedNode] = [
+        #     self.coerce_to_unqualified(start),
+        #     self.coerce_to_unqualified(stop),
+        #     self.coerce_to_unqualified(step),
+        # ]
+        range_args: list[int] = [start, stop, step]
         return UnqualifiedGeneratedCollection(name, column, range_args)
 
 
@@ -857,11 +858,12 @@ def __init__(
 class UnqualifiedGeneratedCollection(UnqualifiedNode):
     """Represents a user-generated collection of values."""
 
-    def __init__(self, name: str, column: list[str], args: list[UnqualifiedNode]):
-        self._parcel: tuple[str, list[str], list[UnqualifiedNode]] = (
+    def __init__(self, name: str, column: list[str], args: list[int]):
+        self._parcel: tuple[str, list[str], list[int]] = (
             name,
             column,
-            [UnqualifiedLiteral(arg, NumericType()) for arg in args],
+            # [UnqualifiedLiteral(arg, NumericType()) for arg in args],
+            args,
         )
 
 
@@ -964,10 +966,10 @@ def display_raw(unqualified: UnqualifiedNode) -> str:
             return result + ")"
         case UnqualifiedGeneratedCollection():
             result = "generated_collection("
-            result += f"name={unqualified._parcel[1]!r}, "
-            result += f"columns=[{', '.join(display_raw(column) for column in unqualified._parcel[2])}],"
+            result += f"name={unqualified._parcel[0]!r}, "
+            result += f"columns=[{', '.join(unqualified._parcel[1])}],"
             result += (
-                f"data=({', '.join(display_raw(arg) for arg in unqualified._parcel[3])})"
+                f"data=({', '.join(unqualified._parcel[2])})"
                 if unqualified._parcel[3:]
                 else ""
             )

pydough/user_collections/__init__.py

@@ -0,0 +1,7 @@
+"""
+Module of PyDough dealing with APIs used for user generated collections.
+"""
+
+__all__ = ["range_collection"]
+
+from .range_collection import range_collection

pydough/user_collections/range_collection.py

@@ -0,0 +1,72 @@
+"""
+Implementation of the `pydough.range_collection` function, which provides
+a way to create a collection of integer ranges over a specified column in PyDough.
+"""
+
+__all__ = ["range_collection"]
+
+from pydough.qdag.collections.range_collection import RangeGeneratedCollection
+
+
+def range_collection(name: str, column: str, *args: int) -> RangeGeneratedCollection:
+    """
+    Creates a collection of integer ranges over a specified column.
+
+    Args:
+        `name` : The name of the collection.
+        `column` : The column to create ranges for.
+        `*args` : Variable length arguments that specify the range parameters.
+        Supported formats:
+            - `range_collection(end)`: generates a range from 0 to `end-1`
+                                        with a step of 1.
+            - `range_collection(start, end)`: generates a range from `start`
+                                        to `end-1` with a step of 1.
+            - `range_collection(start, end, step)`: generates a range from
+                                    `start` to `end-1` with the specified step.
+    Returns:
+        A collection of integer ranges.
+    """
+    if not isinstance(name, str):
+        raise TypeError(f"Expected 'name' to be a string, got {type(name).__name__}")
+    if not isinstance(column, str):
+        raise TypeError(
+            f"Expected 'column' to be a string, got {type(column).__name__}"
+        )
+    if len(args) == 1:
+        end = args[0]
+        start = 0
+        step = 1
+    elif len(args) == 2:
+        start, end = args
+        step = 1
+    elif len(args) == 3:
+        start, end, step = args
+    else:
+        raise ValueError(f"Expected 1 to 3 arguments, got {len(args)}")
+    if not isinstance(start, int):
+        raise TypeError(
+            f"Expected 'start' to be an integer, got {type(start).__name__}"
+        )
+    if not isinstance(end, int):
+        raise TypeError(f"Expected 'end' to be an integer, got {type(end).__name__}")
+    if not isinstance(step, int):
+        raise TypeError(f"Expected 'step' to be an integer, got {type(step).__name__}")
+    if start >= end:
+        raise ValueError(f"Expected 'start' ({start}) to be less than 'end' ({end})")
+    if step == 0:
+        raise ValueError("Expected 'step' to be a non-zero integer")
+    if start < 0:
+        raise ValueError(f"Expected 'start' to be a non-negative integer, got {start}")
+    if end < 0:
+        raise ValueError(f"Expected 'end' to be a non-negative integer, got {end}")
+    # TODO: support negative step values
+    if step <= 0:
+        raise ValueError(f"Expected 'step' to be a positive integer, got {step}")
+
+    return RangeGeneratedCollection(
+        name=name,
+        column_name=column,
+        start=start,
+        end=end,
+        step=step,
+    )

tests/test_pydough_functions/user_collections.py

@@ -0,0 +1,20 @@
+"""
+Various functions containing user generated collections as
+PyDough code snippets for testing purposes.
+"""
+# ruff: noqa
+# mypy: ignore-errors
+# ruff & mypy should not try to typecheck or verify any of this
+
+import pandas as pd
+import datetime
+
+import pydough
+
+
+def simple_range():
+    return pydough.range_collection(
+        "simple_range",
+        "value",
+        10,  # end value
+    ).ORDER_BY(value.ASC())

tests/test_qualification.py

@@ -63,6 +63,9 @@
     impl_tpch_q21,
     impl_tpch_q22,
 )
+from tests.test_pydough_functions.user_collections import (
+    simple_range,
+)
 from tests.testing_utilities import (
     graph_fetcher,
 )
@@ -942,6 +945,12 @@
   """,
             id="simple_cross_6",
         ),
+        pytest.param(
+            simple_range,
+            """
+            """,
+            id="simple_range",
+        ),
     ],
 )
 def test_qualify_node_to_ast_string(