WIP - feat(firestore): add pipeline window functions support

.librarian/generator-input/client-post-processing/firestore-integration.yaml

@@ -92,15 +92,19 @@ replacements:
       )
       from google.cloud.firestore_v1.pipeline_source import PipelineSource
       from google.cloud.firestore_v1.pipeline_types import (
+          DocumentWindowFrame,
           FindNearestOptions,
           Ordering,
           PipelineDataType,
+          RangeWindowFrame,
           SampleOptions,
           SearchOptions,
           TimeGranularity,
           TimePart,
           TimeUnit,
           UnnestOptions,
+          WindowBound,
+          WindowSpec,
       )
       from google.cloud.firestore_v1.query import CollectionGroup, Query
       from google.cloud.firestore_v1.query_profile import (
@@ -198,6 +202,7 @@ replacements:
           "DocumentReference",
           "DocumentSnapshot",
           "DocumentTransform",
+          "DocumentWindowFrame",
           "ExistsOption",
           "ExplainOptions",
           "FieldFilter",
@@ -217,6 +222,7 @@ replacements:
           "PipelineSource",
           "PipelineStream",
           "Query",
+          "RangeWindowFrame",
           "ReadAfterWriteError",
           "SERVER_TIMESTAMP",
           "SampleOptions",
@@ -230,6 +236,8 @@ replacements:
           "types",
           "UnnestOptions",
           "Watch",
+          "WindowBound",
+          "WindowSpec",
           "WriteBatch",
           "WriteOption",
       ]
@@ -278,6 +286,7 @@ replacements:
           DocumentReference,
           DocumentSnapshot,
           DocumentTransform,
+          DocumentWindowFrame,
           ExistsOption,
           ExplainOptions,
           FieldFilter,
@@ -297,6 +306,7 @@ replacements:
           PipelineSource,
           PipelineStream,
           Query,
+          RangeWindowFrame,
           ReadAfterWriteError,
           SampleOptions,
           SearchOptions,
@@ -307,6 +317,8 @@ replacements:
           Transaction,
           UnnestOptions,
           Watch,
+          WindowBound,
+          WindowSpec,
           WriteBatch,
           WriteOption,
           async_transactional,
@@ -345,6 +357,7 @@ replacements:
           "DocumentReference",
           "DocumentSnapshot",
           "DocumentTransform",
+          "DocumentWindowFrame",
           "ExistsOption",
           "ExplainOptions",
           "FieldFilter",
@@ -364,6 +377,7 @@ replacements:
           "PipelineSource",
           "PipelineStream",
           "Query",
+          "RangeWindowFrame",
           "ReadAfterWriteError",
           "SERVER_TIMESTAMP",
           "SampleOptions",
@@ -377,6 +391,8 @@ replacements:
           "types",
           "UnnestOptions",
           "Watch",
+          "WindowBound",
+          "WindowSpec",
           "WriteBatch",
           "WriteOption",
       ]

packages/google-cloud-firestore/google/cloud/firestore/__init__.py

@@ -51,6 +51,7 @@
     DocumentReference,
     DocumentSnapshot,
     DocumentTransform,
+    DocumentWindowFrame,
     ExistsOption,
     ExplainOptions,
     FieldFilter,
@@ -70,6 +71,7 @@
     PipelineSource,
     PipelineStream,
     Query,
+    RangeWindowFrame,
     ReadAfterWriteError,
     SampleOptions,
     SearchOptions,
@@ -80,6 +82,8 @@
     Transaction,
     UnnestOptions,
     Watch,
+    WindowBound,
+    WindowSpec,
     WriteBatch,
     WriteOption,
     async_transactional,
@@ -118,6 +122,7 @@
     "DocumentReference",
     "DocumentSnapshot",
     "DocumentTransform",
+    "DocumentWindowFrame",
     "ExistsOption",
     "ExplainOptions",
     "FieldFilter",
@@ -137,6 +142,7 @@
     "PipelineSource",
     "PipelineStream",
     "Query",
+    "RangeWindowFrame",
     "ReadAfterWriteError",
     "SERVER_TIMESTAMP",
     "SampleOptions",
@@ -150,6 +156,8 @@
     "types",
     "UnnestOptions",
     "Watch",
+    "WindowBound",
+    "WindowSpec",
     "WriteBatch",
     "WriteOption",
 ]

packages/google-cloud-firestore/google/cloud/firestore_v1/__init__.py

@@ -69,15 +69,19 @@
 )
 from google.cloud.firestore_v1.pipeline_source import PipelineSource
 from google.cloud.firestore_v1.pipeline_types import (
+    DocumentWindowFrame,
     FindNearestOptions,
     Ordering,
     PipelineDataType,
+    RangeWindowFrame,
     SampleOptions,
     SearchOptions,
     TimeGranularity,
     TimePart,
     TimeUnit,
     UnnestOptions,
+    WindowBound,
+    WindowSpec,
 )
 from google.cloud.firestore_v1.query import CollectionGroup, Query
 from google.cloud.firestore_v1.query_profile import (
@@ -175,6 +179,7 @@
     "DocumentReference",
     "DocumentSnapshot",
     "DocumentTransform",
+    "DocumentWindowFrame",
     "ExistsOption",
     "ExplainOptions",
     "FieldFilter",
@@ -194,6 +199,7 @@
     "PipelineSource",
     "PipelineStream",
     "Query",
+    "RangeWindowFrame",
     "ReadAfterWriteError",
     "SERVER_TIMESTAMP",
     "SampleOptions",
@@ -207,6 +213,8 @@
     "types",
     "UnnestOptions",
     "Watch",
+    "WindowBound",
+    "WindowSpec",
     "WriteBatch",
     "WriteOption",
 ]

packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py

@@ -27,6 +27,7 @@
     Field,
     FunctionExpression,
     Selectable,
+    WindowFunction,
     _PipelineValueExpression,
 )
 from google.cloud.firestore_v1.types.pipeline import (
@@ -176,6 +177,41 @@ def add_fields(self, *fields: Selectable) -> "_BasePipeline":
         """
         return self._append(stages.AddFields(*fields))
 
+    def add_window_fields(
+        self,
+        *fields: AliasedExpression[WindowFunction],
+        window: types.WindowSpec | None = None,
+    ) -> "_BasePipeline":
+        """
+        Adds new fields computed by window functions over a window of documents.
+
+        .. note::
+            This feature is currently in beta and is subject to change.
+
+        Example:
+            >>> from google.cloud.firestore_v1.pipeline_expressions import Field, Rank
+            >>> from google.cloud.firestore_v1.pipeline_types import WindowSpec
+            >>> pipeline = client.pipeline().collection("books")
+            >>> pipeline = pipeline.add_window_fields(
+            ...     Rank().as_("genre_rank"),
+            ...     Field.of("price").sum().as_("genre_total_price"),
+            ...     window=WindowSpec(
+            ...         partition=["genre"],
+            ...         sort=[Field.of("rating").descending()],
+            ...     ),
+            ... )
+
+        Args:
+            *fields: Aliased `AggregateFunction` or `WindowFunction` expressions defining
+                the window calculations and their output field aliases.
+            window: Optional `WindowSpec` specifying the baseline window configuration
+                for this stage.
+
+        Returns:
+            A new Pipeline object with this stage appended to the stage list.
+        """
+        return self._append(stages.AddWindowFields(*fields, window=window))
+
     def remove_fields(self, *fields: Field | str) -> "_BasePipeline":
         """
         Removes fields from outputs of previous stages.

packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py

@@ -36,6 +36,7 @@
     TimeGranularity,
     TimePart,
     TimeUnit,
+    WindowSpec,
 )
 from google.cloud.firestore_v1.types.document import Pipeline as Pipeline_pb
 from google.cloud.firestore_v1.types.document import Value
@@ -2394,7 +2395,7 @@ def descending(self) -> Ordering:
         """
         return Ordering(self, Ordering.Direction.DESCENDING)
 
-    def as_(self, alias: str) -> "AliasedExpression":
+    def as_(self: T, alias: str) -> "AliasedExpression[T]":
         """Assigns an alias to this expression.
 
         Aliases are useful for renaming fields in the output of a stage or for giving meaningful
@@ -2682,7 +2683,50 @@ def _to_pb(self):
         )
 
 
-class AggregateFunction(FunctionExpression):
+class WindowFunction(FunctionExpression):
+    """A base class for window functions that operate over a window of documents."""
+
+    def over(
+        self,
+        window: WindowSpec,
+    ) -> "WindowFunction":
+        """Configures a window over which this function is evaluated.
+
+        Args:
+            window: A `WindowSpec` specifying the window override for this function.
+
+        Returns:
+            A `WindowFunction` with the specified window override.
+        """
+        return _OverExpression(self, window)
+
+
+class _OverExpression(WindowFunction):
+    """Internal expression binding a WindowFunction to a WindowSpec."""
+
+    def __init__(self, func: WindowFunction, window: WindowSpec):
+        super().__init__(
+            "over",
+            [func],
+            repr_function=lambda expr: f"{expr.params[0]!r}.over({self._window!r})",
+        )
+        self._window = window
+
+    def __eq__(self, other):
+        if not isinstance(other, _OverExpression):
+            return False
+        return super().__eq__(other) and self._window == other._window
+
+    def _to_pb(self) -> Value:
+        return Value(
+            function_value={
+                "name": self.name,
+                "args": [self.params[0]._to_pb(), self._window._to_pb()],
+            }
+        )
+
+
+class AggregateFunction(WindowFunction):
     """A base class for aggregation functions that operate across multiple inputs."""
 
 
@@ -2722,7 +2766,7 @@ def _to_value(field_list: Sequence[Selectable]) -> Value:
         )
 
 
-T = TypeVar("T", bound=Expression)
+T = TypeVar("T", bound=Expression, covariant=True)
 
 
 class AliasedExpression(Selectable, Generic[T]):
@@ -3197,3 +3241,28 @@ def __init__(self):
             [],
             repr_function=FunctionExpression._build_standalone_repr(),
         )
+
+
+class Rank(WindowFunction):
+    """Creates a window function that computes the rank of each document within a window.
+
+    Documents with equal sort values receive the same rank, with gaps in the subsequent
+    ranking numbers.
+
+    Example:
+        >>> # Rank documents by price descending within each category
+        >>> db.pipeline().collection("products").add_window_fields(
+        ...     Rank().as_("price_rank"),
+        ...     window=WindowSpec(
+        ...         partition=["category"],
+        ...         sort=[Field.of("price").descending()],
+        ...     ),
+        ... )
+    """
+
+    def __init__(self):
+        super().__init__(
+            "rank",
+            [],
+            repr_function=FunctionExpression._build_standalone_repr(),
+        )

packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py

@@ -27,6 +27,7 @@
     Expression,
     Field,
     Selectable,
+    WindowFunction,
 )
 from google.cloud.firestore_v1.types.document import Pipeline as Pipeline_pb
 from google.cloud.firestore_v1.types.document import Value
@@ -42,6 +43,7 @@
     SampleOptions,
     SearchOptions,
     UnnestOptions,
+    WindowSpec,
 )
 
 
@@ -86,6 +88,30 @@ def _pb_args(self):
         return [Selectable._to_value(self.fields)]
 
 
+class AddWindowFields(Stage):
+    """Adds new fields computed by window functions over a window of documents."""
+
+    def __init__(
+        self,
+        *args: AliasedExpression[WindowFunction],
+        fields: Sequence[AliasedExpression[WindowFunction]] = (),
+        window: WindowSpec | None = None,
+    ):
+        super().__init__("add_window_fields")
+        if args and fields:
+            raise ValueError(
+                "AddWindowFields stage contains both positional and keyword fields"
+            )
+        self.fields = list(args or fields)
+        self.window: WindowSpec = window if window is not None else WindowSpec()
+
+    def _pb_args(self) -> list[Value]:
+        return [
+            self.window._to_pb(),
+            Selectable._to_value(self.fields),
+        ]
+
+
 class Aggregate(Stage):
     """Performs aggregation operations, optionally grouped."""