fix: window logical plan

daft/__init__.py

@@ -115,6 +115,7 @@
 from daft.sql import sql, sql_expr
 from daft.udf import udf
 from daft.viz import register_viz_hook
+from daft.window import Window
 
 to_struct = Expression.to_struct
 
@@ -134,6 +135,7 @@
     "Session",
     "Table",
     "TimeUnit",
+    "Window",
     "attach_catalog",
     "attach_table",
     "coalesce",

daft/daft/__init__.pyi

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal
 from daft.catalog import Catalog, Table
 from daft.dataframe.display import MermaidOptions
 from daft.execution import physical_plan
+from daft.expressions import Window
 from daft.io.scan import ScanOperator
 from daft.plan_scheduler.physical_plan_scheduler import PartitionT
 from daft.runners.partitioning import PartitionCacheEntry
@@ -69,6 +70,106 @@ class ImageMode(Enum):
         """
         ...
 
+class WindowBoundary:
+    """Represents a window frame boundary in window functions."""
+
+    @staticmethod
+    def UnboundedPreceding() -> WindowBoundary:
+        """Represents UNBOUNDED PRECEDING boundary."""
+        ...
+
+    @staticmethod
+    def UnboundedFollowing() -> WindowBoundary:
+        """Represents UNBOUNDED FOLLOWING boundary."""
+        ...
+
+    @staticmethod
+    def CurrentRow() -> WindowBoundary:
+        """Represents CURRENT ROW boundary."""
+        ...
+
+    @staticmethod
+    def Preceding(n: int) -> WindowBoundary:
+        """Represents N PRECEDING boundary."""
+        ...
+
+    @staticmethod
+    def Following(n: int) -> WindowBoundary:
+        """Represents N FOLLOWING boundary."""
+        ...
+
+class WindowFrameType:
+    """Represents the type of window frame (ROWS or RANGE)."""
+
+    @staticmethod
+    def Rows() -> WindowFrameType:
+        """Row-based window frame."""
+        ...
+
+    @staticmethod
+    def Range() -> WindowFrameType:
+        """Range-based window frame."""
+        ...
+
+class WindowFrame:
+    """Represents a window frame specification."""
+
+    def __init__(
+        self,
+        frame_type: WindowFrameType,
+        start: WindowBoundary,
+        end: WindowBoundary,
+    ) -> None:
+        """Create a new window frame specification.
+
+        Args:
+            frame_type: Type of window frame (ROWS or RANGE)
+            start: Start boundary of window frame
+            end: End boundary of window frame
+        """
+        ...
+
+class WindowSpec:
+    """Represents a window specification for window functions."""
+
+    @staticmethod
+    def new() -> WindowSpec:
+        """Create a new empty window specification."""
+        ...
+
+    def with_partition_by(self, exprs: list[PyExpr]) -> WindowSpec:
+        """Set the partition by expressions.
+
+        Args:
+            exprs: List of expressions to partition by
+        """
+        ...
+
+    def with_order_by(self, exprs: list[PyExpr], ascending: list[bool]) -> WindowSpec:
+        """Set the order by expressions.
+
+        Args:
+            exprs: List of expressions to order by
+            ascending: List of booleans indicating sort order for each expression
+        """
+        ...
+
+    def with_frame(self, frame: WindowFrame) -> WindowSpec:
+        """Set the window frame specification.
+
+        Args:
+            frame: Window frame specification
+        """
+        ...
+
+    def with_min_periods(self, min_periods: int) -> WindowSpec:
+        """Set the minimum number of rows required to compute a result.
+
+        Args:
+            min_periods: Minimum number of rows required
+        """
+        ...
+
 class ImageFormat(Enum):
     """Supported image formats for Daft's image I/O."""
 
@@ -954,6 +1055,7 @@ class PyExpr:
     def agg_list(self) -> PyExpr: ...
     def agg_set(self) -> PyExpr: ...
     def agg_concat(self) -> PyExpr: ...
+    def over(self, window_spec: Window) -> PyExpr: ...
     def __add__(self, other: PyExpr) -> PyExpr: ...
     def __sub__(self, other: PyExpr) -> PyExpr: ...
     def __mul__(self, other: PyExpr) -> PyExpr: ...

daft/expressions/expressions.py

@@ -54,6 +54,31 @@
 if TYPE_CHECKING:
     from daft.io import IOConfig
     from daft.udf import BoundUDFArgs, InitArgsType, UninitializedUdf
+    from daft.window import Window
+
+    # Type hints for window functions
+    def _window_func(expr: _PyExpr) -> _PyExpr: ...
+
+    native.rank = _window_func
+    native.dense_rank = _window_func
+    native.row_number = _window_func
+    native.percent_rank = _window_func
+    native.first_value = _window_func
+    native.last_value = _window_func
+
+    def _ntile(expr: _PyExpr, n: int) -> _PyExpr: ...
+
+    native.ntile = _ntile
+
+    def _nth_value(expr: _PyExpr, n: int) -> _PyExpr: ...
+
+    native.nth_value = _nth_value
+
+    def _lag_lead(expr: _PyExpr, offset: int, default: _PyExpr | None) -> _PyExpr: ...
+
+    native.lag = _lag_lead
+    native.lead = _lag_lead
+
 # This allows Sphinx to correctly work against our "namespaced" accessor functions by overriding @property to
 # return a class instance of the namespace instead of a property object.
 elif os.getenv("DAFT_SPHINX_BUILD") == "1":
@@ -1468,6 +1493,126 @@
     def _initialize_udfs(self) -> Expression:
         return Expression._from_pyexpr(initialize_udfs(self._expr))
 
+    def over(self, window: Window) -> Expression:
+        """Apply this expression as a window function over the specified window.
+
+        Args:
+            window: Window specification defining partitioning and ordering
+
+        Returns:
+            Expression: A new expression representing the window function result
+        """
+        return Expression._from_pyexpr(self._expr.over(window._spec))
+
+    def rank(self) -> Expression:
+        """Compute rank within window partition.
+
+        Ranks are consecutive integers starting from 1, with gaps for ties.
+        For example, if two rows tie for rank 2, the next rank will be 4.
+
+        Returns:
+            Expression: Expression containing rank values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def dense_rank(self) -> Expression:
+        """Compute dense rank within window partition.
+
+        Dense ranks are consecutive integers starting from 1, without gaps for ties.
+        For example, if two rows tie for rank 2, the next rank will be 3.
+
+        Returns:
+            Expression: Expression containing dense rank values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def row_number(self) -> Expression:
+        """Compute row number within window partition.
+
+        Row numbers are consecutive integers starting from 1, assigned to each row
+        in the partition based on the window's ordering.
+
+        Returns:
+            Expression: Expression containing row numbers
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def percent_rank(self) -> Expression:
+        """Compute percent rank within window partition.
+
+        Percent rank is (rank - 1) / (partition_rows - 1), ranging from 0 to 1.
+        Returns NULL if partition has only one row.
+
+        Returns:
+            Expression: Expression containing percent rank values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def ntile(self, n: int) -> Expression:
+        """Divide rows in partition into n buckets numbered from 1 to n.
+
+        Buckets are assigned as evenly as possible, with remaining rows
+        distributed one per bucket starting from bucket 1.
+
+        Args:
+            n: Number of buckets to divide rows into
+
+        Returns:
+            Expression: Expression containing bucket numbers
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def lag(self, offset: int = 1, default: Any = None) -> Expression:
+        """Access value from previous row in partition.
+
+        Args:
+            offset: Number of rows to look back (default: 1)
+            default: Value to return if no previous row exists (default: None)
+
+        Returns:
+            Expression: Expression containing lagged values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def lead(self, offset: int = 1, default: Any = None) -> Expression:
+        """Access value from following row in partition.
+
+        Args:
+            offset: Number of rows to look ahead (default: 1)
+            default: Value to return if no following row exists (default: None)
+
+        Returns:
+            Expression: Expression containing leading values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def first_value(self) -> Expression:
+        """Get first value in window frame.
+
+        Returns:
+            Expression: Expression containing first values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def last_value(self) -> Expression:
+        """Get last value in window frame.
+
+        Returns:
+            Expression: Expression containing last values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
+    def nth_value(self, n: int) -> Expression:
+        """Get nth value in window frame.
+
+        Args:
+            n: Position of value to get (1-based)
+
+        Returns:
+            Expression: Expression containing nth values
+        """
+        raise NotImplementedError("Window functions are not yet implemented")
+
 
 SomeExpressionNamespace = TypeVar("SomeExpressionNamespace", bound="ExpressionNamespace")