From fa96508a330221c3d55f66f0f7e588dee40f8476 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Mon, 2 Nov 2026 20:58:46 +0000 Subject: [PATCH 1/9] feat(firestore): add pipeline window functions support Add WindowSpec, DocumentWindowFrame, RangeWindowFrame, and WindowBound types in pipeline_types, the AddWindowFields stage and add_window_fields() pipeline builder method, and the Rank() window function and Expression.over() frame override method. Include unit tests and pipeline_e2e acceptance tests in window.yaml. --- .../firestore-integration.yaml | 16 + .../google/cloud/firestore/__init__.py | 8 + .../google/cloud/firestore_v1/__init__.py | 8 + .../cloud/firestore_v1/base_pipeline.py | 36 ++ .../firestore_v1/pipeline_expressions.py | 49 ++ .../cloud/firestore_v1/pipeline_stages.py | 26 + .../cloud/firestore_v1/pipeline_types.py | 206 +++++++ .../tests/system/pipeline_e2e/window.yaml | 506 ++++++++++++++++++ .../tests/unit/v1/test_async_pipeline.py | 17 + .../tests/unit/v1/test_pipeline.py | 17 + .../unit/v1/test_pipeline_expressions.py | 46 ++ .../tests/unit/v1/test_pipeline_stages.py | 192 +++++++ 12 files changed, 1127 insertions(+) create mode 100644 packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml 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() From a2ccd2f57f93f18a1ddc0976efcd1a675e3954f4 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 3 Nov 2026 18:46:08 +0000 Subject: [PATCH 2/9] test(firestore): resolve pipeline_types classes in acceptance test parser Look up pipeline option and window frame classes (such as DocumentWindowFrame and RangeWindowFrame) directly in pipeline_types within _parse_expressions and _is_stage_string. --- .../tests/system/test_pipeline_acceptance.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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..23480dea510e 100644 --- a/packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py +++ b/packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py @@ -31,6 +31,7 @@ from google.cloud.firestore_v1 import pipeline_expressions 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 import pipeline_types from google.cloud.firestore_v1.vector import Vector FIRESTORE_PROJECT = os.environ.get("GCLOUD_PROJECT") @@ -318,7 +319,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 +400,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) ) From 623d5ac89f68aa49a7af3900e6b52db07df1a17d Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 3 Nov 2026 20:56:54 +0000 Subject: [PATCH 3/9] refactor(firestore): move over() to WindowFunction Define over() on WindowFunction instead of Expression and make AggregateFunction inherit from WindowFunction. This restricts over() to window and aggregate functions and simplifies AddWindowFields type annotations. Also sort imports in test_pipeline_acceptance.py. --- .../cloud/firestore_v1/base_pipeline.py | 2 +- .../firestore_v1/pipeline_expressions.py | 37 +++++++++---------- .../cloud/firestore_v1/pipeline_stages.py | 4 +- .../tests/system/pipeline_e2e/window.yaml | 4 +- .../tests/system/test_pipeline_acceptance.py | 3 +- .../unit/v1/test_pipeline_expressions.py | 11 +++++- 6 files changed, 33 insertions(+), 28 deletions(-) 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 afe8b19c69ad..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 @@ -179,7 +179,7 @@ def add_fields(self, *fields: Selectable) -> "_BasePipeline": def add_window_fields( self, - *fields: AliasedExpression[AggregateFunction | WindowFunction], + *fields: AliasedExpression[WindowFunction], window: types.WindowSpec | None = None, ) -> "_BasePipeline": """ 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 e43526cde0a4..ed58dbb78172 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 @@ -2575,21 +2575,6 @@ 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.""" @@ -2698,14 +2683,26 @@ def _to_pb(self): ) -class AggregateFunction(FunctionExpression): - """A base class for aggregation functions that operate across multiple inputs.""" +class WindowFunction(FunctionExpression): + """A base class for window functions that operate over a window of documents.""" - pass + 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. -class WindowFunction(FunctionExpression): - """A base class for window functions that operate over a window of documents.""" + Returns: + A `WindowFunction` with the specified window override. + """ + return WindowFunction("over", [self, window]) # type: ignore[list-item] + + +class AggregateFunction(WindowFunction): + """A base class for aggregation functions that operate across multiple inputs.""" pass 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 3003c178a376..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 @@ -93,8 +93,8 @@ class AddWindowFields(Stage): def __init__( self, - *args: AliasedExpression[AggregateFunction | WindowFunction], - fields: Sequence[AliasedExpression[AggregateFunction | WindowFunction]] = (), + *args: AliasedExpression[WindowFunction], + fields: Sequence[AliasedExpression[WindowFunction]] = (), window: WindowSpec | None = None, ): super().__init__("add_window_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 index 07e011a5404c..2a46c2eccf60 100644 --- a/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml +++ b/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml @@ -143,7 +143,7 @@ tests: - Rank - "rating_rank" - AliasedExpression: - - FunctionExpression.over: + - WindowFunction.over: - FunctionExpression.sum: - Field: published - WindowSpec: @@ -153,7 +153,7 @@ tests: following: 0.15 - "nearby_published_sum" - AliasedExpression: - - FunctionExpression.over: + - WindowFunction.over: - Count - WindowSpec: documents: 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 23480dea510e..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,10 +28,9 @@ 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 import pipeline_types from google.cloud.firestore_v1.vector import Vector FIRESTORE_PROJECT = os.environ.get("GCLOUD_PROJECT") 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 4d4c2008e9d5..3ca223ae2206 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,6 +22,7 @@ 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, @@ -2269,17 +2270,25 @@ def test_array_index_of_all(self): class TestWindowExpressions: def test_over_and_rank(self): + assert issubclass(AggregateFunction, WindowFunction) + assert not hasattr(Expression, "over") + assert not hasattr(Field.of("rating"), "over") + agg = Field.of("rating").sum() + assert isinstance(agg, AggregateFunction) + assert isinstance(agg, WindowFunction) ws = WindowSpec( partition="genre", sort=Field.of("published").ascending(), ) win_fn = agg.over(ws) assert isinstance(win_fn, WindowFunction) + assert not isinstance(win_fn, AggregateFunction) 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) + assert win_fn == WindowFunction.over(agg, ws) + assert win_fn == AggregateFunction.over(agg, ws) pb = win_fn._to_pb() assert pb.function_value.name == "over" From 311f7b103258a3d4d264f4a1481f64d2cb2d9ba4 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 3 Nov 2026 21:45:22 +0000 Subject: [PATCH 4/9] refactor(firestore): use dataclass and _WindowFrame for window frames Convert DocumentWindowFrame and RangeWindowFrame to dataclasses inheriting from a shared _WindowFrame base class to eliminate duplicated __init__, __repr__, __eq__, and _to_pb boilerplate. --- .../cloud/firestore_v1/pipeline_types.py | 129 +++++------------- 1 file changed, 35 insertions(+), 94 deletions(-) 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 2c707c1cfd48..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 @@ -336,7 +337,23 @@ class WindowBound(str, Enum): UNBOUNDED = "unbounded" -class DocumentWindowFrame: +@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: @@ -348,49 +365,16 @@ class DocumentWindowFrame: 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}) + preceding: Optional[WindowBound | str | int | pipeline_expressions.Expression] = ( + None + ) + following: Optional[WindowBound | str | int | pipeline_expressions.Expression] = ( + None + ) -class RangeWindowFrame: +@dataclass +class RangeWindowFrame(_WindowFrame): """Defines a value-based or time-range window frame relative to the current document's sort value. Args: @@ -400,58 +384,15 @@ class RangeWindowFrame: `"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}) + 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: From 456647affc8738c7edf44de80c43ed783ba172bc Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 4 Nov 2026 01:32:00 +0000 Subject: [PATCH 5/9] refactor(firestore): add _OverExpression and covariant AliasedExpression Introduce _OverExpression as a WindowFunction subclass to bind a WindowFunction to a WindowSpec without passing WindowSpec into FunctionExpression.params or suppressing mypy list-item checks. Make AliasedExpression type variable T covariant and preserve T in Expression.as_() so AliasedExpression[AggregateFunction] is assignable to AliasedExpression[WindowFunction]. --- .../firestore_v1/pipeline_expressions.py | 31 +++++++++++++++++-- .../unit/v1/test_pipeline_expressions.py | 5 ++- 2 files changed, 32 insertions(+), 4 deletions(-) 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 ed58dbb78172..83ad29807cf5 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 @@ -2395,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 @@ -2698,7 +2698,32 @@ def over( Returns: A `WindowFunction` with the specified window override. """ - return WindowFunction("over", [self, window]) # type: ignore[list-item] + 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): @@ -2743,7 +2768,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]): 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 3ca223ae2206..64665d3bc98c 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 @@ -2285,10 +2285,13 @@ def test_over_and_rank(self): assert isinstance(win_fn, WindowFunction) assert not isinstance(win_fn, AggregateFunction) assert win_fn.name == "over" - assert win_fn.params == [agg, ws] + assert win_fn.params == [agg] + assert win_fn._window == ws assert repr(win_fn) == f"Field.of('rating').sum().over({ws!r})" assert win_fn == WindowFunction.over(agg, ws) assert win_fn == AggregateFunction.over(agg, ws) + assert win_fn != agg.over(WindowSpec()) + assert win_fn != agg pb = win_fn._to_pb() assert pb.function_value.name == "over" From d1ab1a4a0126d65bc3a8485da3754b57029b8b04 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 4 Nov 2026 02:17:00 +0000 Subject: [PATCH 6/9] test(firestore): simplify window unit tests and align with backend semantics Split test_over_and_rank into test_over and test_rank, remove redundant dataclass boilerplate assertions, and apply .over() frame overrides to aggregate functions rather than Rank() to match Firestore backend rules. --- .../unit/v1/test_pipeline_expressions.py | 18 +++------- .../tests/unit/v1/test_pipeline_stages.py | 33 ++++--------------- 2 files changed, 12 insertions(+), 39 deletions(-) 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 64665d3bc98c..9dab49ff8fea 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 @@ -32,6 +32,7 @@ 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 @@ -2269,7 +2270,7 @@ def test_array_index_of_all(self): class TestWindowExpressions: - def test_over_and_rank(self): + def test_over(self): assert issubclass(AggregateFunction, WindowFunction) assert not hasattr(Expression, "over") assert not hasattr(Field.of("rating"), "over") @@ -2278,8 +2279,7 @@ def test_over_and_rank(self): assert isinstance(agg, AggregateFunction) assert isinstance(agg, WindowFunction) ws = WindowSpec( - partition="genre", - sort=Field.of("published").ascending(), + documents=DocumentWindowFrame(-1, WindowBound.CURRENT), ) win_fn = agg.over(ws) assert isinstance(win_fn, WindowFunction) @@ -2288,8 +2288,7 @@ def test_over_and_rank(self): assert win_fn.params == [agg] assert win_fn._window == ws assert repr(win_fn) == f"Field.of('rating').sum().over({ws!r})" - assert win_fn == WindowFunction.over(agg, ws) - assert win_fn == AggregateFunction.over(agg, ws) + assert win_fn == agg.over(ws) assert win_fn != agg.over(WindowSpec()) assert win_fn != agg @@ -2306,17 +2305,10 @@ def test_over_and_rank(self): assert count_pb.function_value.name == "over" assert count_pb.function_value.args[0] == expr.Count()._to_pb() - # Rank + def test_rank(self): 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 252be3f47e06..3be1bbe7b75e 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 @@ -1079,13 +1079,6 @@ 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) @@ -1110,13 +1103,6 @@ def test_range_window_frame(self): 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) @@ -1178,12 +1164,7 @@ def test_window_spec_init_and_coercions(self): 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" - ) + assert pb.map_value.fields["documents"] == spec_doc.documents._to_pb() # RangeWindowFrame spec_range = WindowSpec(range=RangeWindowFrame(-5, 5, "day")) @@ -1232,16 +1213,16 @@ def test_ctor_errors(self): 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") + expr1 = ( + Field.of("rating") + .sum() + .over(WindowSpec(documents=DocumentWindowFrame(-1, WindowBound.CURRENT))) + .as_("sum_rating") ) + expr2 = Rank().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 ( From 11ff92a3fd040ba37279813a6d63f66b0af9858c Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 4 Nov 2026 02:23:52 +0000 Subject: [PATCH 7/9] chore(firestore): remove redundant pass statement from AggregateFunction --- .../google/cloud/firestore_v1/pipeline_expressions.py | 2 -- 1 file changed, 2 deletions(-) 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 83ad29807cf5..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 @@ -2729,8 +2729,6 @@ def _to_pb(self) -> Value: class AggregateFunction(WindowFunction): """A base class for aggregation functions that operate across multiple inputs.""" - pass - class Selectable(Expression): """Base class for expressions that can be selected or aliased in projection stages.""" From f26431dd89b0f8a8372a5e15401de0763806d604 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 4 Nov 2026 19:33:24 +0000 Subject: [PATCH 8/9] test(firestore): decompose and tighten pipeline window unit tests Split multi-scenario window expression, type, and stage tests into focused single-purpose test methods. Remove dataclass field-echo and private attribute assertions, cover TimeGranularity in RangeWindowFrame, and import WindowSpec from pipeline_types. --- .../tests/unit/v1/test_async_pipeline.py | 3 +- .../tests/unit/v1/test_pipeline.py | 3 +- .../unit/v1/test_pipeline_expressions.py | 41 ++-- .../tests/unit/v1/test_pipeline_stages.py | 227 ++++++++---------- 4 files changed, 121 insertions(+), 153 deletions(-) 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 c4bfb4a7b5ea..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 @@ -457,7 +458,7 @@ def test_async_pipeline_aggregate_with_groups(): def test_async_pipeline_add_window_fields_with_window(): start_ppl = _make_async_pipeline() - window = stages.WindowSpec(partition="genre", sort=Field.of("rating").descending()) + 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 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 863ab47c68b9..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 @@ -448,7 +449,7 @@ def test_pipeline_aggregate_with_groups(): def test_pipeline_add_window_fields_with_window(): start_ppl = _make_pipeline() - window = stages.WindowSpec(partition="genre", sort=Field.of("rating").descending()) + 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 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 9dab49ff8fea..6db5dd8d1bfe 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 @@ -2270,45 +2270,34 @@ def test_array_index_of_all(self): class TestWindowExpressions: - def test_over(self): + def test_window_function_hierarchy(self): assert issubclass(AggregateFunction, WindowFunction) - assert not hasattr(Expression, "over") + assert not issubclass(Expression, WindowFunction) assert not hasattr(Field.of("rating"), "over") + def test_over_repr_and_equality(self): agg = Field.of("rating").sum() - assert isinstance(agg, AggregateFunction) - assert isinstance(agg, WindowFunction) - ws = WindowSpec( - documents=DocumentWindowFrame(-1, WindowBound.CURRENT), - ) - win_fn = agg.over(ws) + window = WindowSpec(partition="genre") + win_fn = agg.over(window) assert isinstance(win_fn, WindowFunction) assert not isinstance(win_fn, AggregateFunction) - assert win_fn.name == "over" - assert win_fn.params == [agg] - assert win_fn._window == ws - assert repr(win_fn) == f"Field.of('rating').sum().over({ws!r})" - assert win_fn == agg.over(ws) + assert ( + repr(win_fn) + == "Field.of('rating').sum().over(WindowSpec(partition=[Field.of('genre')]))" + ) + assert win_fn == agg.over(WindowSpec(partition="genre")) assert win_fn != agg.over(WindowSpec()) assert win_fn != agg - pb = win_fn._to_pb() + 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 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() + assert list(pb.function_value.args) == [agg._to_pb(), window._to_pb()] def test_rank(self): 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": []}) 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 3be1bbe7b75e..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 @@ -30,6 +30,7 @@ from google.cloud.firestore_v1.pipeline_types import ( DocumentWindowFrame, RangeWindowFrame, + TimeGranularity, TimeUnit, WindowBound, WindowSpec, @@ -1075,165 +1076,141 @@ def test_to_pb_with_fields(self): class TestWindowTypes: - def test_document_window_frame(self): - frame = DocumentWindowFrame(preceding=1, following=WindowBound.CURRENT) - assert frame.preceding == 1 - assert frame.following == WindowBound.CURRENT - - 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 - - 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" - ) + 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"), + } - 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" - ) + @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), + } - empty_frame = RangeWindowFrame() - assert empty_frame._to_pb() == Value(map_value={"fields": {}}) + 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_init_and_coercions(self): - # Empty WindowSpec + def test_window_spec_repr_and_equality(self): 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( + spec = WindowSpec( partition="genre", - sort=sort_order, + sort=Field.of("rating").descending(), documents=DocumentWindowFrame(-1, WindowBound.CURRENT), + range=RangeWindowFrame(-5, 5, "day"), ) - 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 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_doc == WindowSpec( + assert spec == WindowSpec( partition=[Field.of("genre")], sort=[Field.of("rating").descending()], documents=DocumentWindowFrame(-1, WindowBound.CURRENT), + range=RangeWindowFrame(-5, 5, "day"), ) - 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"] == spec_doc.documents._to_pb() - - # 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'))" + 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 ) - range_pb = spec_range._to_pb() - assert range_pb.map_value.fields["range"] == spec_range.range._to_pb() + 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_and_keyword(self): + 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" - # 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): + 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_and_to_pb(self): - expr1 = ( - Field.of("rating") - .sum() - .over(WindowSpec(documents=DocumentWindowFrame(-1, WindowBound.CURRENT))) - .as_("sum_rating") + def test_repr(self): + stage = self._make_one( + Rank().as_("pub_rank"), window=WindowSpec(partition="genre") ) - expr2 = Rank().as_("pub_rank") - window = WindowSpec( - partition="genre", - sort=Field.of("rating").descending(), - ) - stage = self._make_one(expr1, expr2, window=window) - assert ( - repr(stage) - == f"AddWindowFields(fields={[expr1, expr2]!r}, window={window!r})" + assert repr(stage) == ( + "AddWindowFields(" + "fields=[Rank().as_('pub_rank')], " + "window=WindowSpec(partition=[Field.of('genre')])" + ")" ) - pb = stage._to_pb() + 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 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() + assert pb.args[1].map_value.fields == { + "sum_rating": expr1.expr._to_pb(), + "pub_rank": expr2.expr._to_pb(), + } From 8abe03c87d007b3210fa2ec1b0329c2f3eb25eee Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 4 Nov 2026 23:42:23 +0000 Subject: [PATCH 9/9] test(firestore): split rank and over() frame window system tests Separate rank with ties (sorted by rating ASC to exercise tie ranking) from accumulator-level over() frames (sorted by unique published ASC to avoid positional tie ambiguity and year-sum math) in tests/system/pipeline_e2e/window.yaml. Remove the redundant expression-partition window system test whose permutations are already covered by unit tests, and update test_over_repr_and_equality to pass a valid frame-only WindowSpec. --- .../tests/system/pipeline_e2e/window.yaml | 348 ++++++++---------- .../unit/v1/test_pipeline_expressions.py | 8 +- 2 files changed, 155 insertions(+), 201 deletions(-) diff --git a/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml b/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml index 2a46c2eccf60..66669aa169bc 100644 --- a/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml +++ b/packages/google-cloud-firestore/tests/system/pipeline_e2e/window.yaml @@ -124,7 +124,7 @@ tests: title: fieldReferenceValue: title name: select - - description: "testAddWindowFields - rank and accumulator-level over() frames" + - description: "testAddWindowFields - rank with ties" pipeline: - Collection: books - Where: @@ -142,56 +142,159 @@ tests: - 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: - - FunctionExpression.sum: - - Field: published + - Count - WindowSpec: range: RangeWindowFrame: - preceding: 0.15 - following: 0.15 - - "nearby_published_sum" + preceding: 20 + following: current + - "books_within_20_years" - AliasedExpression: - WindowFunction.over: - - Count + - FunctionExpression.array_agg: + - Field: title - WindowSpec: documents: DocumentWindowFrame: - preceding: -1 - following: 2 - - "look_ahead_count" + preceding: 1 + following: 1 + - "nearby_titles" - Sort: - Ordering: - Field: published - ASCENDING - Select: - title - - rating - - rating_rank - - nearby_published_sum - - look_ahead_count + - published + - books_within_20_years + - nearby_titles assert_results: - title: "1984" - rating: 4.2 - rating_rank: 2 - nearby_published_sum: 5913 - look_ahead_count: 2 + published: 1949 + books_within_20_years: 1 + nearby_titles: + - "1984" + - "Dune" - title: "Dune" - rating: 4.6 - rating_rank: 4 - nearby_published_sum: 1965 - look_ahead_count: 0 + 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" - rating: 4.2 - rating_rank: 2 - nearby_published_sum: 5913 - look_ahead_count: 1 + 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" - rating: 4.1 - rating_rank: 1 - nearby_published_sum: 5913 - look_ahead_count: 2 + published: 1985 + books_within_20_years: 3 + nearby_titles: + - "The Hitchhiker's Guide to the Galaxy" + - "The Handmaid's Tale" assert_proto: pipeline: stages: @@ -220,44 +323,41 @@ tests: direction: stringValue: ascending expression: - fieldReferenceValue: rating + fieldReferenceValue: published - mapValue: fields: - look_ahead_count: + books_within_20_years: functionValue: args: - functionValue: name: count - mapValue: fields: - documents: + range: mapValue: fields: following: - integerValue: '2' + stringValue: current preceding: - integerValue: '-1' + integerValue: '20' name: over - nearby_published_sum: + nearby_titles: functionValue: args: - functionValue: args: - - fieldReferenceValue: published - name: sum + - fieldReferenceValue: title + name: array_agg - mapValue: fields: - range: + documents: mapValue: fields: following: - doubleValue: 0.15 + integerValue: '1' preceding: - doubleValue: 0.15 + integerValue: '1' name: over - rating_rank: - functionValue: - name: rank name: add_window_fields - args: - mapValue: @@ -270,14 +370,12 @@ tests: - 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 + books_within_20_years: + fieldReferenceValue: books_within_20_years + nearby_titles: + fieldReferenceValue: nearby_titles + published: + fieldReferenceValue: published title: fieldReferenceValue: title name: select @@ -343,151 +441,6 @@ tests: 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 @@ -503,4 +456,3 @@ tests: - Count - "window_count" assert_error: ".*range.*" - 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 6db5dd8d1bfe..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 @@ -2277,15 +2277,17 @@ def test_window_function_hierarchy(self): def test_over_repr_and_equality(self): agg = Field.of("rating").sum() - window = WindowSpec(partition="genre") + 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(partition=[Field.of('genre')]))" + == "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(partition="genre")) assert win_fn != agg.over(WindowSpec()) assert win_fn != agg