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..afe8b19c69ad 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[AggregateFunction | 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..e43526cde0a4 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 @@ -2574,6 +2575,21 @@ def is_type( "is_type", [self, self._cast_to_expr_or_convert_to_constant(type_val)] ) + @expose_as_static + 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 WindowFunction("over", [self, window]) # type: ignore[list-item] + class Constant(Expression, Generic[CONSTANT_TYPE]): """Represents a constant literal value in an expression.""" @@ -2685,6 +2701,14 @@ def _to_pb(self): class AggregateFunction(FunctionExpression): """A base class for aggregation functions that operate across multiple inputs.""" + pass + + +class WindowFunction(FunctionExpression): + """A base class for window functions that operate over a window of documents.""" + + pass + class Selectable(Expression): """Base class for expressions that can be selected or aliased in projection stages.""" @@ -3197,3 +3221,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..3003c178a376 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[AggregateFunction | WindowFunction], + fields: Sequence[AliasedExpression[AggregateFunction | 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..2c707c1cfd48 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 @@ -327,3 +327,209 @@ 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" + + +class DocumentWindowFrame: + """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. + """ + + def __init__( + self, + preceding: Optional[ + WindowBound | str | int | pipeline_expressions.Expression + ] = None, + following: Optional[ + WindowBound | str | int | pipeline_expressions.Expression + ] = None, + ): + self.preceding = preceding + self.following = following + + 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.preceding is not None: + args.append(f"preceding={self.preceding!r}") + if self.following is not None: + args.append(f"following={self.following!r}") + return f"{self.__class__.__name__}({', '.join(args)})" + + def _to_pb(self) -> Value: + fields: dict[str, Value] = {} + if self.preceding is not None: + fields["preceding"] = ( + pipeline_expressions.Expression._cast_to_expr_or_convert_to_constant( + self.preceding + )._to_pb() + ) + if self.following is not None: + fields["following"] = ( + pipeline_expressions.Expression._cast_to_expr_or_convert_to_constant( + self.following + )._to_pb() + ) + return Value(map_value={"fields": fields}) + + +class RangeWindowFrame: + """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`). + """ + + def __init__( + self, + 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, + ): + self.preceding = preceding + self.following = following + self.unit = unit + + 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.preceding is not None: + args.append(f"preceding={self.preceding!r}") + if self.following is not None: + args.append(f"following={self.following!r}") + if self.unit is not None: + args.append(f"unit={self.unit!r}") + return f"{self.__class__.__name__}({', '.join(args)})" + + def _to_pb(self) -> Value: + fields: dict[str, Value] = {} + if self.preceding is not None: + fields["preceding"] = ( + pipeline_expressions.Expression._cast_to_expr_or_convert_to_constant( + self.preceding + )._to_pb() + ) + if self.following is not None: + fields["following"] = ( + pipeline_expressions.Expression._cast_to_expr_or_convert_to_constant( + self.following + )._to_pb() + ) + if self.unit is not None: + fields["unit"] = ( + pipeline_expressions.Expression._cast_to_expr_or_convert_to_constant( + self.unit + )._to_pb() + ) + return Value(map_value={"fields": fields}) + + +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..07e011a5404c --- /dev/null +++ b/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml @@ -0,0 +1,506 @@ +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 and accumulator-level over() frames" + 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" + - AliasedExpression: + - FunctionExpression.over: + - FunctionExpression.sum: + - Field: published + - WindowSpec: + range: + RangeWindowFrame: + preceding: 0.15 + following: 0.15 + - "nearby_published_sum" + - AliasedExpression: + - FunctionExpression.over: + - Count + - WindowSpec: + documents: + DocumentWindowFrame: + preceding: -1 + following: 2 + - "look_ahead_count" + - Sort: + - Ordering: + - Field: published + - ASCENDING + - Select: + - title + - rating + - rating_rank + - nearby_published_sum + - look_ahead_count + assert_results: + - title: "1984" + rating: 4.2 + rating_rank: 2 + nearby_published_sum: 5913 + look_ahead_count: 2 + - title: "Dune" + rating: 4.6 + rating_rank: 4 + nearby_published_sum: 1965 + look_ahead_count: 0 + - title: "The Hitchhiker's Guide to the Galaxy" + rating: 4.2 + rating_rank: 2 + nearby_published_sum: 5913 + look_ahead_count: 1 + - title: "The Handmaid's Tale" + rating: 4.1 + rating_rank: 1 + nearby_published_sum: 5913 + look_ahead_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: + sort: + arrayValue: + values: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: rating + - mapValue: + fields: + look_ahead_count: + functionValue: + args: + - functionValue: + name: count + - mapValue: + fields: + documents: + mapValue: + fields: + following: + integerValue: '2' + preceding: + integerValue: '-1' + name: over + nearby_published_sum: + functionValue: + args: + - functionValue: + args: + - fieldReferenceValue: published + name: sum + - mapValue: + fields: + range: + mapValue: + fields: + following: + doubleValue: 0.15 + preceding: + doubleValue: 0.15 + name: over + rating_rank: + functionValue: + name: rank + name: add_window_fields + - args: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: published + name: sort + - args: + - mapValue: + fields: + look_ahead_count: + fieldReferenceValue: look_ahead_count + nearby_published_sum: + fieldReferenceValue: nearby_published_sum + rating: + fieldReferenceValue: rating + rating_rank: + fieldReferenceValue: rating_rank + 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 - expression partition, multiple sorts, array_agg, and nested alias" + pipeline: + - Collection: books + - Where: + - FunctionExpression.equal_any: + - Field: genre + - [Science Fiction, Dystopian] + - AddWindowFields: + window: + WindowSpec: + partition: + - FunctionExpression.to_upper: + - Field: genre + sort: + - Ordering: + - Field: published + - ASCENDING + - Ordering: + - Field: title + - ASCENDING + documents: + DocumentWindowFrame: + preceding: unbounded + following: + Constant: 0 + fields: + - AliasedExpression: + - FunctionExpression.sum: + - Field: published + - "stats.running_published" + - AliasedExpression: + - FunctionExpression.array_agg: + - Field: title + - "running_titles" + - Sort: + - Ordering: + - Field: published + - ASCENDING + - Select: + - title + - stats + - running_titles + assert_results: + - title: "1984" + stats: + running_published: 1949 + running_titles: + - "1984" + - title: "Dune" + stats: + running_published: 1965 + running_titles: + - "Dune" + - title: "The Hitchhiker's Guide to the Galaxy" + stats: + running_published: 3944 + running_titles: + - "Dune" + - "The Hitchhiker's Guide to the Galaxy" + - title: "The Handmaid's Tale" + stats: + running_published: 3934 + running_titles: + - "1984" + - "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: + documents: + mapValue: + fields: + following: + integerValue: '0' + preceding: + stringValue: unbounded + partition: + arrayValue: + values: + - functionValue: + args: + - fieldReferenceValue: genre + name: to_upper + sort: + arrayValue: + values: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: published + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: title + - mapValue: + fields: + running_titles: + functionValue: + args: + - fieldReferenceValue: title + name: array_agg + stats.running_published: + functionValue: + args: + - fieldReferenceValue: published + name: sum + name: add_window_fields + - args: + - mapValue: + fields: + direction: + stringValue: ascending + expression: + fieldReferenceValue: published + name: sort + - args: + - mapValue: + fields: + running_titles: + fieldReferenceValue: running_titles + stats: + fieldReferenceValue: stats + title: + fieldReferenceValue: title + 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/unit/v1/test_async_pipeline.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py index 5f9e7f2749ab..c4bfb4a7b5ea 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 @@ -426,6 +426,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 +455,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 = stages.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..863ab47c68b9 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py @@ -417,6 +417,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 +446,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 = stages.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..4d4c2008e9d5 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 @@ -27,6 +27,9 @@ Expression, Field, Ordering, + Rank, + WindowFunction, + WindowSpec, ) from google.cloud.firestore_v1.types import document as document_pb from google.cloud.firestore_v1.types import query as query_pb @@ -2262,3 +2265,46 @@ 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_over_and_rank(self): + agg = Field.of("rating").sum() + ws = WindowSpec( + partition="genre", + sort=Field.of("published").ascending(), + ) + win_fn = agg.over(ws) + assert isinstance(win_fn, WindowFunction) + assert win_fn.name == "over" + assert win_fn.params == [agg, ws] + assert repr(win_fn) == f"Field.of('rating').sum().over({ws!r})" + assert win_fn == Expression.over(agg, ws) + + pb = win_fn._to_pb() + assert pb.function_value.name == "over" + assert len(pb.function_value.args) == 2 + assert pb.function_value.args[0] == agg._to_pb() + assert pb.function_value.args[1] == ws._to_pb() + + # Also test with Count() (0-arg AggregateFunction) + count_win = expr.Count().over(WindowSpec()) + assert repr(count_win) == "Count().over(WindowSpec())" + count_pb = count_win._to_pb() + assert count_pb.function_value.name == "over" + assert count_pb.function_value.args[0] == expr.Count()._to_pb() + + # Rank + rank = Rank() + assert isinstance(rank, WindowFunction) + assert rank.name == "rank" + assert rank.params == [] + assert repr(rank) == "Rank()" + assert rank._to_pb() == Value(function_value={"name": "rank", "args": []}) + + rank_over = rank.over(ws) + assert isinstance(rank_over, WindowFunction) + assert repr(rank_over) == f"Rank().over({ws!r})" + assert rank_over._to_pb() == Value( + function_value={"name": "over", "args": [rank._to_pb(), ws._to_pb()]} + ) 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..252be3f47e06 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,14 @@ DocumentMatches, Field, Ordering, + Rank, +) +from google.cloud.firestore_v1.pipeline_types import ( + DocumentWindowFrame, + RangeWindowFrame, + TimeUnit, + WindowBound, + WindowSpec, ) from google.cloud.firestore_v1.types.document import Value from google.cloud.firestore_v1.vector import Vector @@ -1064,3 +1072,187 @@ 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(self): + frame = DocumentWindowFrame(preceding=1, following=WindowBound.CURRENT) + assert frame.preceding == 1 + assert frame.following == WindowBound.CURRENT + assert ( + repr(frame) + == "DocumentWindowFrame(preceding=1, following=)" + ) + assert frame == DocumentWindowFrame(1, WindowBound.CURRENT) + assert frame != DocumentWindowFrame(2, WindowBound.CURRENT) + assert frame != object() + + pb = frame._to_pb() + assert pb.map_value.fields["preceding"] == Value(integer_value=1) + assert pb.map_value.fields["following"] == Value(string_value="current") + + expr_frame = DocumentWindowFrame( + preceding=Field.of("offset"), following="unbounded" + ) + expr_pb = expr_frame._to_pb() + assert expr_pb.map_value.fields["preceding"] == Value( + field_reference_value="offset" + ) + assert expr_pb.map_value.fields["following"] == Value(string_value="unbounded") + + empty_frame = DocumentWindowFrame() + assert empty_frame._to_pb() == Value(map_value={"fields": {}}) + + def test_range_window_frame(self): + frame = RangeWindowFrame( + preceding=10.5, following=WindowBound.UNBOUNDED, unit=TimeUnit.DAY + ) + assert frame.preceding == 10.5 + assert frame.following == WindowBound.UNBOUNDED + assert frame.unit == TimeUnit.DAY + assert ( + repr(frame) + == "RangeWindowFrame(preceding=10.5, following=, unit=)" + ) + assert frame == RangeWindowFrame(10.5, WindowBound.UNBOUNDED, TimeUnit.DAY) + assert frame != RangeWindowFrame(10.5, WindowBound.UNBOUNDED, "hour") + assert frame != object() + + pb = frame._to_pb() + assert pb.map_value.fields["preceding"] == Value(double_value=10.5) + assert pb.map_value.fields["following"] == Value(string_value="unbounded") + assert pb.map_value.fields["unit"] == Value(string_value="day") + + str_unit_frame = RangeWindowFrame(preceding=5, following=0, unit="hour") + assert str_unit_frame._to_pb().map_value.fields["unit"] == Value( + string_value="hour" + ) + + expr_unit_frame = RangeWindowFrame( + preceding=5, following=0, unit=Constant.of("minute") + ) + assert expr_unit_frame._to_pb().map_value.fields["unit"] == Value( + string_value="minute" + ) + + empty_frame = RangeWindowFrame() + assert empty_frame._to_pb() == Value(map_value={"fields": {}}) + + def test_window_spec_init_and_coercions(self): + # Empty WindowSpec + empty = WindowSpec() + assert empty.partition == [] + assert empty.sort == [] + assert empty.documents is None + assert empty.range is None + assert repr(empty) == "WindowSpec()" + assert empty._to_pb() == Value(map_value={"fields": {}}) + + # Single-item partition and sort coercion + DocumentWindowFrame + sort_order = Field.of("rating").descending() + spec_doc = WindowSpec( + partition="genre", + sort=sort_order, + documents=DocumentWindowFrame(-1, WindowBound.CURRENT), + ) + assert spec_doc.partition == [Field.of("genre")] + assert spec_doc.sort == [sort_order] + assert spec_doc.documents == DocumentWindowFrame(-1, WindowBound.CURRENT) + assert spec_doc.range is None + assert ( + repr(spec_doc) + == "WindowSpec(partition=[Field.of('genre')], sort=[Field.of('rating').descending()], documents=DocumentWindowFrame(preceding=-1, following=))" + ) + assert spec_doc == WindowSpec( + partition=[Field.of("genre")], + sort=[Field.of("rating").descending()], + documents=DocumentWindowFrame(-1, WindowBound.CURRENT), + ) + assert spec_doc != empty + assert spec_doc != object() + + pb = spec_doc._to_pb() + assert list(pb.map_value.fields["partition"].array_value.values) == [ + Value(field_reference_value="genre") + ] + assert list(pb.map_value.fields["sort"].array_value.values) == [ + Field.of("rating").descending()._to_pb() + ] + assert pb.map_value.fields["documents"].map_value.fields["preceding"] == Value( + integer_value=-1 + ) + assert pb.map_value.fields["documents"].map_value.fields["following"] == Value( + string_value="current" + ) + + # RangeWindowFrame + spec_range = WindowSpec(range=RangeWindowFrame(-5, 5, "day")) + assert spec_range.range == RangeWindowFrame(-5, 5, "day") + assert ( + repr(spec_range) + == "WindowSpec(range=RangeWindowFrame(preceding=-5, following=5, unit='day'))" + ) + range_pb = spec_range._to_pb() + assert range_pb.map_value.fields["range"] == spec_range.range._to_pb() + + +class TestAddWindowFields: + def _make_one(self, *args, **kwargs): + return stages.AddWindowFields(*args, **kwargs) + + def test_ctor_positional_and_keyword(self): + expr1 = Field.of("rating").sum().as_("sum_rating") + expr2 = Rank().as_("rank_val") + + # Positional + stage1 = self._make_one(expr1, expr2) + assert stage1.fields == [expr1, expr2] + assert stage1.window == WindowSpec() + assert stage1.name == "add_window_fields" + + # Keyword fields + window + stage2 = self._make_one( + fields=[expr1], + window=WindowSpec( + partition=["genre"], sort=[Field.of("rating").descending()] + ), + ) + assert stage2.fields == [expr1] + assert stage2.window == WindowSpec( + partition=["genre"], sort=[Field.of("rating").descending()] + ) + + def test_ctor_errors(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_and_to_pb(self): + expr1 = Field.of("rating").sum().as_("sum_rating") + expr2 = ( + Rank() + .over(WindowSpec(sort=Field.of("published").ascending())) + .as_("pub_rank") + ) + window = WindowSpec( + partition="genre", + sort=Field.of("rating").descending(), + documents=DocumentWindowFrame(-1, WindowBound.CURRENT), + ) + stage = self._make_one(expr1, expr2, window=window) + assert ( + repr(stage) + == f"AddWindowFields(fields={[expr1, expr2]!r}, window={window!r})" + ) + + pb = stage._to_pb() + assert pb.name == "add_window_fields" + assert len(pb.options) == 0 + assert len(pb.args) == 2 + assert pb.args[0] == window._to_pb() + assert pb.args[1].map_value.fields["sum_rating"] == expr1.expr._to_pb() + assert pb.args[1].map_value.fields["pub_rank"] == expr2.expr._to_pb()