diff --git a/.librarian/generator-input/client-post-processing/firestore-integration.yaml b/.librarian/generator-input/client-post-processing/firestore-integration.yaml index d31022a3b362..18af4ae9800b 100644 --- a/.librarian/generator-input/client-post-processing/firestore-integration.yaml +++ b/.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", ] diff --git a/packages/google-cloud-firestore/google/cloud/firestore/__init__.py b/packages/google-cloud-firestore/google/cloud/firestore/__init__.py index 44a9651d6c02..a45b406145db 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore/__init__.py +++ b/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", ] diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/__init__.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/__init__.py index 4656b5b5441f..a5a2757afc57 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/__init__.py +++ b/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", ] diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py index e0699b59004f..3c5ab6910dfa 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py +++ b/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. diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py index 5f52d92fd3bd..2e5b2ae5dc35 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py +++ b/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(), + ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py index e72523dc9dae..6a6fb0e13b85 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py +++ b/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.""" diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_types.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_types.py index 3e9918140e1f..c543e41bfddb 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_types.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_types.py @@ -15,6 +15,7 @@ from __future__ import annotations +from dataclasses import dataclass, fields from enum import Enum from typing import Optional, Sequence @@ -327,3 +328,149 @@ def __init__(self, index_field: pipeline_expressions.Field | str): def __repr__(self): return f"{self.__class__.__name__}(index_field={self.index_field.path!r})" + + +class WindowBound(str, Enum): + """Enumeration of symbolic window frame boundary keywords.""" + + CURRENT = "current" + UNBOUNDED = "unbounded" + + +@dataclass +class _WindowFrame: + """Base class for window frame specifications.""" + + def _to_pb(self) -> Value: + pb_fields = { + f.name: pipeline_expressions.Expression._cast_to_expr_or_convert_to_constant( + getattr(self, f.name) + )._to_pb() + for f in fields(self) + if getattr(self, f.name) is not None + } + return Value(map_value={"fields": pb_fields}) + + +@dataclass +class DocumentWindowFrame(_WindowFrame): + """Defines a document-count based window frame relative to the current document. + + Args: + preceding: The lower bound of the frame (an integer offset, `WindowBound`, or `Expression`). + A positive integer includes up to that many documents before the current document. + A negative integer indicates a boundary following the current document. + following: The upper bound of the frame (an integer offset, `WindowBound`, or `Expression`). + A positive integer includes up to that many documents after the current document. + A negative integer indicates a boundary preceding the current document. + """ + + preceding: Optional[WindowBound | str | int | pipeline_expressions.Expression] = ( + None + ) + following: Optional[WindowBound | str | int | pipeline_expressions.Expression] = ( + None + ) + + +@dataclass +class RangeWindowFrame(_WindowFrame): + """Defines a value-based or time-range window frame relative to the current document's sort value. + + Args: + preceding: The lower bound of the frame (a numeric offset, `WindowBound`, or `Expression`). + following: The upper bound of the frame (a numeric offset, `WindowBound`, or `Expression`). + unit: Optional duration unit for time-range frames (e.g., `TimeUnit.DAY`, `TimeGranularity.MONTH`, + `"day"`, or an `Expression`). + """ + + preceding: Optional[ + WindowBound | str | int | float | pipeline_expressions.Expression + ] = None + following: Optional[ + WindowBound | str | int | float | pipeline_expressions.Expression + ] = None + unit: Optional[ + TimeUnit | TimeGranularity | str | pipeline_expressions.Expression + ] = None + + +class WindowSpec: + """Configuration for a window in an `AddWindowFields` pipeline stage or `.over()` expression. + + Args: + partition: Field paths or expressions used to partition documents into independent windows. + sort: Orderings specifying how documents are sorted within each partition. + documents: A `DocumentWindowFrame` defining a document-count window frame. + Mutually exclusive with `range`. + range: A `RangeWindowFrame` defining a value-based or time-range window frame. + Mutually exclusive with `documents`. + """ + + def __init__( + self, + *, + partition: Optional[ + Sequence[pipeline_expressions.Expression | str] + | pipeline_expressions.Expression + | str + ] = None, + sort: Optional[Sequence[Ordering] | Ordering] = None, + documents: Optional[DocumentWindowFrame] = None, + range: Optional[RangeWindowFrame] = None, + ): + if partition is None: + self.partition: list[pipeline_expressions.Expression] = [] + else: + partition_seq = ( + [partition] + if isinstance(partition, (str, pipeline_expressions.Expression)) + else list(partition) + ) + self.partition = [ + p + if isinstance(p, pipeline_expressions.Expression) + else pipeline_expressions.Field.of(p) + for p in partition_seq + ] + + if sort is None: + self.sort: list[Ordering] = [] + else: + self.sort = [sort] if isinstance(sort, Ordering) else list(sort) + + self.documents: Optional[DocumentWindowFrame] = documents + self.range: Optional[RangeWindowFrame] = range + + def __eq__(self, other): + if not isinstance(other, type(self)): + return False + return self._to_pb() == other._to_pb() + + def __repr__(self): + args = [] + if self.partition: + args.append(f"partition={self.partition!r}") + if self.sort: + args.append(f"sort={self.sort!r}") + if self.documents is not None: + args.append(f"documents={self.documents!r}") + if self.range is not None: + args.append(f"range={self.range!r}") + return f"{self.__class__.__name__}({', '.join(args)})" + + def _to_pb(self) -> Value: + fields: dict[str, Value] = {} + if self.partition: + fields["partition"] = Value( + array_value={"values": [p._to_pb() for p in self.partition]} + ) + if self.sort: + fields["sort"] = Value( + array_value={"values": [s._to_pb() for s in self.sort]} + ) + if self.documents is not None: + fields["documents"] = self.documents._to_pb() + if self.range is not None: + fields["range"] = self.range._to_pb() + return Value(map_value={"fields": fields}) diff --git a/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml b/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml new file mode 100644 index 000000000000..66669aa169bc --- /dev/null +++ b/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml @@ -0,0 +1,458 @@ +tests: + - description: "testAddWindowFields - partition, sort, documents frame, and multiple accumulators" + pipeline: + - Collection: books + - Where: + - FunctionExpression.equal_any: + - Field: genre + - [Science Fiction, Dystopian] + - AddWindowFields: + window: + WindowSpec: + partition: + - Field: genre + sort: + - Ordering: + - Field: published + - ASCENDING + documents: + DocumentWindowFrame: + preceding: unbounded + following: current + fields: + - AliasedExpression: + - FunctionExpression.sum: + - Field: rating + - "running_rating" + - AliasedExpression: + - Count + - "running_count" + - Sort: + - Ordering: + - Field: published + - ASCENDING + - Select: + - title + - genre + - running_rating + - running_count + assert_results: + - title: "1984" + genre: "Dystopian" + running_rating: 4.2 + running_count: 1 + - title: "Dune" + genre: "Science Fiction" + running_rating: 4.6 + running_count: 1 + - title: "The Hitchhiker's Guide to the Galaxy" + genre: "Science Fiction" + running_rating: 8.8 + running_count: 2 + - title: "The Handmaid's Tale" + genre: "Dystopian" + running_rating: 8.3 + running_count: 2 + assert_proto: + pipeline: + stages: + - args: + - referenceValue: /books + name: collection + - args: + - functionValue: + args: + - fieldReferenceValue: genre + - functionValue: + args: + - stringValue: Science Fiction + - stringValue: Dystopian + name: array + name: equal_any + name: where + - args: + - mapValue: + fields: + documents: + mapValue: + fields: + following: + stringValue: current + preceding: + stringValue: unbounded + partition: + arrayValue: + values: + - fieldReferenceValue: genre + sort: + arrayValue: + values: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: published + - mapValue: + fields: + running_count: + functionValue: + name: count + running_rating: + functionValue: + args: + - fieldReferenceValue: rating + name: sum + name: add_window_fields + - args: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: published + name: sort + - args: + - mapValue: + fields: + genre: + fieldReferenceValue: genre + running_count: + fieldReferenceValue: running_count + running_rating: + fieldReferenceValue: running_rating + title: + fieldReferenceValue: title + name: select + - description: "testAddWindowFields - rank with ties" + pipeline: + - Collection: books + - Where: + - FunctionExpression.equal_any: + - Field: genre + - [Science Fiction, Dystopian] + - AddWindowFields: + window: + WindowSpec: + sort: + - Ordering: + - Field: rating + - ASCENDING + fields: + - AliasedExpression: + - Rank + - "rating_rank" + - Sort: + - Ordering: + - Field: rating + - ASCENDING + - Ordering: + - Field: title + - ASCENDING + - Select: + - title + - rating + - rating_rank + assert_results: + - title: "The Handmaid's Tale" + rating: 4.1 + rating_rank: 1 + - title: "1984" + rating: 4.2 + rating_rank: 2 + - title: "The Hitchhiker's Guide to the Galaxy" + rating: 4.2 + rating_rank: 2 + - title: "Dune" + rating: 4.6 + rating_rank: 4 + assert_proto: + pipeline: + stages: + - args: + - referenceValue: /books + name: collection + - args: + - functionValue: + args: + - fieldReferenceValue: genre + - functionValue: + args: + - stringValue: Science Fiction + - stringValue: Dystopian + name: array + name: equal_any + name: where + - args: + - mapValue: + fields: + sort: + arrayValue: + values: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: rating + - mapValue: + fields: + rating_rank: + functionValue: + name: rank + name: add_window_fields + - args: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: rating + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: title + name: sort + - args: + - mapValue: + fields: + rating: + fieldReferenceValue: rating + rating_rank: + fieldReferenceValue: rating_rank + title: + fieldReferenceValue: title + name: select + - description: "testAddWindowFields - accumulator-level over() frames" + pipeline: + - Collection: books + - Where: + - FunctionExpression.equal_any: + - Field: genre + - [Science Fiction, Dystopian] + - AddWindowFields: + window: + WindowSpec: + sort: + - Ordering: + - Field: published + - ASCENDING + fields: + - AliasedExpression: + - WindowFunction.over: + - Count + - WindowSpec: + range: + RangeWindowFrame: + preceding: 20 + following: current + - "books_within_20_years" + - AliasedExpression: + - WindowFunction.over: + - FunctionExpression.array_agg: + - Field: title + - WindowSpec: + documents: + DocumentWindowFrame: + preceding: 1 + following: 1 + - "nearby_titles" + - Sort: + - Ordering: + - Field: published + - ASCENDING + - Select: + - title + - published + - books_within_20_years + - nearby_titles + assert_results: + - title: "1984" + published: 1949 + books_within_20_years: 1 + nearby_titles: + - "1984" + - "Dune" + - title: "Dune" + published: 1965 + books_within_20_years: 2 + nearby_titles: + - "1984" + - "Dune" + - "The Hitchhiker's Guide to the Galaxy" + - title: "The Hitchhiker's Guide to the Galaxy" + published: 1979 + books_within_20_years: 2 + nearby_titles: + - "Dune" + - "The Hitchhiker's Guide to the Galaxy" + - "The Handmaid's Tale" + - title: "The Handmaid's Tale" + published: 1985 + books_within_20_years: 3 + nearby_titles: + - "The Hitchhiker's Guide to the Galaxy" + - "The Handmaid's Tale" + assert_proto: + pipeline: + stages: + - args: + - referenceValue: /books + name: collection + - args: + - functionValue: + args: + - fieldReferenceValue: genre + - functionValue: + args: + - stringValue: Science Fiction + - stringValue: Dystopian + name: array + name: equal_any + name: where + - args: + - mapValue: + fields: + sort: + arrayValue: + values: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: published + - mapValue: + fields: + books_within_20_years: + functionValue: + args: + - functionValue: + name: count + - mapValue: + fields: + range: + mapValue: + fields: + following: + stringValue: current + preceding: + integerValue: '20' + name: over + nearby_titles: + functionValue: + args: + - functionValue: + args: + - fieldReferenceValue: title + name: array_agg + - mapValue: + fields: + documents: + mapValue: + fields: + following: + integerValue: '1' + preceding: + integerValue: '1' + name: over + name: add_window_fields + - args: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: published + name: sort + - args: + - mapValue: + fields: + books_within_20_years: + fieldReferenceValue: books_within_20_years + nearby_titles: + fieldReferenceValue: nearby_titles + published: + fieldReferenceValue: published + title: + fieldReferenceValue: title + name: select + - description: "testAddWindowFields - range frame with time unit" + pipeline: + - Collection: timestamps + - AddWindowFields: + window: + WindowSpec: + sort: + - Ordering: + - Field: time + - ASCENDING + range: + RangeWindowFrame: + preceding: 12 + following: current + unit: hour + fields: + - AliasedExpression: + - Count + - "window_count" + - Select: + - window_count + assert_results: + - window_count: 1 + assert_proto: + pipeline: + stages: + - args: + - referenceValue: /timestamps + name: collection + - args: + - mapValue: + fields: + range: + mapValue: + fields: + following: + stringValue: current + preceding: + integerValue: '12' + unit: + stringValue: hour + sort: + arrayValue: + values: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: time + - mapValue: + fields: + window_count: + functionValue: + name: count + name: add_window_fields + - args: + - mapValue: + fields: + window_count: + fieldReferenceValue: window_count + name: select + - description: "testAddWindowFields - error on range frame without sort" + pipeline: + - Collection: books + - AddWindowFields: + window: + WindowSpec: + range: + RangeWindowFrame: + preceding: unbounded + following: current + fields: + - AliasedExpression: + - Count + - "window_count" + assert_error: ".*range.*" diff --git a/packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py b/packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py index 149182b13694..ab9f03a652f5 100644 --- a/packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py +++ b/packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py @@ -28,7 +28,7 @@ from test__helpers import FIRESTORE_EMULATOR, FIRESTORE_ENTERPRISE_DB, system_test_lock from google.cloud.firestore import AsyncClient, Client, GeoPoint -from google.cloud.firestore_v1 import pipeline_expressions +from google.cloud.firestore_v1 import pipeline_expressions, pipeline_types from google.cloud.firestore_v1 import pipeline_expressions as expr from google.cloud.firestore_v1 import pipeline_stages as stages from google.cloud.firestore_v1.vector import Vector @@ -318,7 +318,7 @@ def _parse_expressions(client, yaml_element: Any): elif len(yaml_element) == 1 and _is_stage_string(next(iter(yaml_element))): # build pipeline stage if possible (eg, for SampleOptions) cls_str = next(iter(yaml_element)) - cls = getattr(stages, cls_str) + cls = getattr(pipeline_types, cls_str) yaml_args = yaml_element[cls_str] return _apply_yaml_args_to_callable(cls, client, yaml_args) elif len(yaml_element) == 1 and list(yaml_element)[0] == "Pipeline": @@ -399,7 +399,7 @@ def _is_stage_string(yaml_str): return ( isinstance(yaml_str, str) and yaml_str[0].isupper() - and hasattr(stages, yaml_str) + and hasattr(pipeline_types, yaml_str) ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py index 5f9e7f2749ab..629613e1f326 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py @@ -17,6 +17,7 @@ from google.cloud.firestore_v1 import pipeline_stages as stages from google.cloud.firestore_v1.pipeline_expressions import Field +from google.cloud.firestore_v1.pipeline_types import WindowSpec from tests.unit.v1._test_helpers import make_async_client @@ -426,6 +427,11 @@ async def test_async_pipeline_stream_stream_equivalence(): ("offset", (1,), stages.Offset), ("limit", (1,), stages.Limit), ("aggregate", (Field.of("n").as_("alias"),), stages.Aggregate), + ( + "add_window_fields", + (Field.of("n").sum().as_("alias"),), + stages.AddWindowFields, + ), ("distinct", ("field_name",), stages.Distinct), ("distinct", (Field.of("n"), "second"), stages.Distinct), ], @@ -450,6 +456,18 @@ def test_async_pipeline_aggregate_with_groups(): assert list(result_ppl.stages[0].accumulators) == [Field.of("title")] +def test_async_pipeline_add_window_fields_with_window(): + start_ppl = _make_async_pipeline() + window = WindowSpec(partition="genre", sort=Field.of("rating").descending()) + expr = Field.of("rating").sum().as_("sum_rating") + result_ppl = start_ppl.add_window_fields(expr, window=window) + assert len(start_ppl.stages) == 0 + assert len(result_ppl.stages) == 1 + assert isinstance(result_ppl.stages[0], stages.AddWindowFields) + assert result_ppl.stages[0].fields == [expr] + assert result_ppl.stages[0].window == window + + def test_async_pipeline_raw_stage_with_options(): from google.cloud.firestore_v1.pipeline_stages import RawStage diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py index fa2f19ae109a..9617a2a070e7 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py @@ -17,6 +17,7 @@ from google.cloud.firestore_v1 import pipeline_stages as stages from google.cloud.firestore_v1.pipeline_expressions import Field +from google.cloud.firestore_v1.pipeline_types import WindowSpec from tests.unit.v1._test_helpers import make_client @@ -417,6 +418,11 @@ def test_pipeline_execute_stream_equivalence(): ("offset", (1,), stages.Offset), ("limit", (1,), stages.Limit), ("aggregate", (Field.of("n").as_("alias"),), stages.Aggregate), + ( + "add_window_fields", + (Field.of("n").sum().as_("alias"),), + stages.AddWindowFields, + ), ("distinct", ("field_name",), stages.Distinct), ("distinct", (Field.of("n"), "second"), stages.Distinct), ], @@ -441,6 +447,18 @@ def test_pipeline_aggregate_with_groups(): assert list(result_ppl.stages[0].accumulators) == [Field.of("title")] +def test_pipeline_add_window_fields_with_window(): + start_ppl = _make_pipeline() + window = WindowSpec(partition="genre", sort=Field.of("rating").descending()) + expr = Field.of("rating").sum().as_("sum_rating") + result_ppl = start_ppl.add_window_fields(expr, window=window) + assert len(start_ppl.stages) == 0 + assert len(result_ppl.stages) == 1 + assert isinstance(result_ppl.stages[0], stages.AddWindowFields) + assert result_ppl.stages[0].fields == [expr] + assert result_ppl.stages[0].window == window + + def test_pipeline_raw_stage_with_options(): from google.cloud.firestore_v1.pipeline_stages import RawStage diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py index 0c0f29df54b3..50f415e81981 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py @@ -22,12 +22,17 @@ from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1._helpers import GeoPoint from google.cloud.firestore_v1.pipeline_expressions import ( + AggregateFunction, BooleanExpression, Constant, Expression, Field, Ordering, + Rank, + WindowFunction, + WindowSpec, ) +from google.cloud.firestore_v1.pipeline_types import DocumentWindowFrame, WindowBound from google.cloud.firestore_v1.types import document as document_pb from google.cloud.firestore_v1.types import query as query_pb from google.cloud.firestore_v1.types.document import Value @@ -2262,3 +2267,39 @@ def test_array_index_of_all(self): assert repr(instance) == "Value.array_index_of_all(Constant.of('comedy'))" infix_instance = arg1.array_index_of_all(value) assert infix_instance == instance + + +class TestWindowExpressions: + def test_window_function_hierarchy(self): + assert issubclass(AggregateFunction, WindowFunction) + assert not issubclass(Expression, WindowFunction) + assert not hasattr(Field.of("rating"), "over") + + def test_over_repr_and_equality(self): + agg = Field.of("rating").sum() + window = WindowSpec(documents=DocumentWindowFrame(1, WindowBound.CURRENT)) + win_fn = agg.over(window) + assert isinstance(win_fn, WindowFunction) + assert not isinstance(win_fn, AggregateFunction) + assert ( + repr(win_fn) + == "Field.of('rating').sum().over(WindowSpec(documents=DocumentWindowFrame(preceding=1, following=)))" + ) + assert win_fn == agg.over( + WindowSpec(documents=DocumentWindowFrame(1, WindowBound.CURRENT)) + ) + assert win_fn != agg.over(WindowSpec()) + assert win_fn != agg + + def test_over_to_pb(self): + agg = Field.of("rating").sum() + window = WindowSpec(documents=DocumentWindowFrame(-1, WindowBound.CURRENT)) + pb = agg.over(window)._to_pb() + assert pb.function_value.name == "over" + assert list(pb.function_value.args) == [agg._to_pb(), window._to_pb()] + + def test_rank(self): + rank = Rank() + assert isinstance(rank, WindowFunction) + assert repr(rank) == "Rank()" + assert rank._to_pb() == Value(function_value={"name": "rank", "args": []}) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py index c926285f462c..12bd771ec0d3 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py @@ -25,6 +25,15 @@ DocumentMatches, Field, Ordering, + Rank, +) +from google.cloud.firestore_v1.pipeline_types import ( + DocumentWindowFrame, + RangeWindowFrame, + TimeGranularity, + TimeUnit, + WindowBound, + WindowSpec, ) from google.cloud.firestore_v1.types.document import Value from google.cloud.firestore_v1.vector import Vector @@ -1064,3 +1073,144 @@ def test_to_pb_with_fields(self): assert "score" in result.args[0].map_value.fields assert "status" in result.args[0].map_value.fields assert len(result.options) == 0 + + +class TestWindowTypes: + def test_document_window_frame_to_pb(self): + assert DocumentWindowFrame()._to_pb().map_value.fields == {} + pb = DocumentWindowFrame(1, WindowBound.CURRENT)._to_pb() + assert pb.map_value.fields == { + "preceding": Value(integer_value=1), + "following": Value(string_value="current"), + } + expr_pb = DocumentWindowFrame(Field.of("offset"), "unbounded")._to_pb() + assert expr_pb.map_value.fields == { + "preceding": Value(field_reference_value="offset"), + "following": Value(string_value="unbounded"), + } + + @pytest.mark.parametrize( + "unit,expected_unit", + [ + (TimeUnit.DAY, "day"), + (TimeGranularity.MONTH, "month"), + ("hour", "hour"), + (Constant.of("minute"), "minute"), + ], + ) + def test_range_window_frame_to_pb(self, unit, expected_unit): + assert RangeWindowFrame()._to_pb().map_value.fields == {} + pb = RangeWindowFrame(10.5, WindowBound.UNBOUNDED, unit)._to_pb() + assert pb.map_value.fields == { + "preceding": Value(double_value=10.5), + "following": Value(string_value="unbounded"), + "unit": Value(string_value=expected_unit), + } + + def test_window_spec_coercions(self): + order = Field.of("rating").descending() + assert WindowSpec().partition == [] + assert WindowSpec().sort == [] + assert WindowSpec(partition="genre").partition == [Field.of("genre")] + assert WindowSpec(partition=Field.of("genre")).partition == [Field.of("genre")] + assert WindowSpec(partition=["genre", Field.of("author")]).partition == [ + Field.of("genre"), + Field.of("author"), + ] + assert WindowSpec(sort=order).sort == [order] + assert WindowSpec(sort=[order]).sort == [order] + + def test_window_spec_repr_and_equality(self): + empty = WindowSpec() + assert repr(empty) == "WindowSpec()" + spec = WindowSpec( + partition="genre", + sort=Field.of("rating").descending(), + documents=DocumentWindowFrame(-1, WindowBound.CURRENT), + range=RangeWindowFrame(-5, 5, "day"), + ) + assert repr(spec) == ( + "WindowSpec(" + "partition=[Field.of('genre')], " + "sort=[Field.of('rating').descending()], " + "documents=DocumentWindowFrame(preceding=-1, following=), " + "range=RangeWindowFrame(preceding=-5, following=5, unit='day')" + ")" + ) + assert spec == WindowSpec( + partition=[Field.of("genre")], + sort=[Field.of("rating").descending()], + documents=DocumentWindowFrame(-1, WindowBound.CURRENT), + range=RangeWindowFrame(-5, 5, "day"), + ) + assert spec != empty + assert spec != object() + + def test_window_spec_to_pb(self): + assert WindowSpec()._to_pb().map_value.fields == {} + order = Field.of("rating").descending() + doc_frame = DocumentWindowFrame(-1, WindowBound.CURRENT) + range_frame = RangeWindowFrame(-5, 5, "day") + spec = WindowSpec( + partition="genre", sort=order, documents=doc_frame, range=range_frame + ) + assert spec._to_pb().map_value.fields == { + "partition": Value( + array_value={"values": [Value(field_reference_value="genre")]} + ), + "sort": Value(array_value={"values": [order._to_pb()]}), + "documents": doc_frame._to_pb(), + "range": range_frame._to_pb(), + } + + +class TestAddWindowFields: + def _make_one(self, *args, **kwargs): + return stages.AddWindowFields(*args, **kwargs) + + def test_ctor_positional(self): + expr1 = Field.of("rating").sum().as_("sum_rating") + expr2 = Rank().as_("rank_val") + stage = self._make_one(expr1, expr2) + assert stage.fields == [expr1, expr2] + assert stage.window == WindowSpec() + assert stage.name == "add_window_fields" + + def test_ctor_keyword(self): + expr1 = Field.of("rating").sum().as_("sum_rating") + window = WindowSpec(partition=["genre"], sort=[Field.of("rating").descending()]) + stage = self._make_one(fields=[expr1], window=window) + assert stage.fields == [expr1] + assert stage.window == window + + def test_ctor_combined_raises(self): + expr1 = Field.of("rating").sum().as_("sum_rating") + with pytest.raises( + ValueError, + match="AddWindowFields stage contains both positional and keyword fields", + ): + self._make_one(expr1, fields=[expr1]) + + def test_repr(self): + stage = self._make_one( + Rank().as_("pub_rank"), window=WindowSpec(partition="genre") + ) + assert repr(stage) == ( + "AddWindowFields(" + "fields=[Rank().as_('pub_rank')], " + "window=WindowSpec(partition=[Field.of('genre')])" + ")" + ) + + def test_to_pb(self): + expr1 = Field.of("rating").sum().as_("sum_rating") + expr2 = Rank().as_("pub_rank") + window = WindowSpec(partition="genre") + pb = self._make_one(expr1, expr2, window=window)._to_pb() + assert pb.name == "add_window_fields" + assert len(pb.options) == 0 + assert pb.args[0] == window._to_pb() + assert pb.args[1].map_value.fields == { + "sum_rating": expr1.expr._to_pb(), + "pub_rank": expr2.expr._to_pb(), + }