Merge branch 'main' into shuowei-execution-history

Author: shuoweil

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

@@ -394,6 +394,37 @@ def sort(self, *orders: stages.Ordering) -> "_BasePipeline":
         """
         return self._append(stages.Sort(*orders))
 
+    def search(
+        self, query_or_options: str | BooleanExpression | stages.SearchOptions
+    ) -> "_BasePipeline":
+        """
+        Adds a search stage to the pipeline.
+
+        This stage filters documents based on the provided query expression.
+
+        Example:
+            >>> from google.cloud.firestore_v1.pipeline_stages import SearchOptions
+            >>> from google.cloud.firestore_v1.pipeline_expressions import And, DocumentMatches, Field, GeoPoint
+            >>> # Search for restaurants matching either "waffles" or "pancakes" near a location
+            >>> pipeline = client.pipeline().collection("restaurants").search(
+            ...     SearchOptions(
+            ...         query=And(
+            ...             DocumentMatches("waffles OR pancakes"),
+            ...             Field.of("location").geo_distance(GeoPoint(38.9, -107.0)).less_than(1000)
+            ...         ),
+            ...         sort=Score().descending()
+            ...     )
+            ... )
+
+        Args:
+            options: Either a string or expression representing the search query, or
+                A `SearchOptions` instance configuring the search.
+
+        Returns:
+            A new Pipeline object with this stage appended to the stage list
+        """
+        return self._append(stages.Search(query_or_options))
+
     def sample(self, limit_or_options: int | stages.SampleOptions) -> "_BasePipeline":
         """
         Performs a pseudo-random sampling of the documents from the previous stage.

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

@@ -730,6 +730,61 @@ def less_than_or_equal(
             [self, self._cast_to_expr_or_convert_to_constant(other)],
         )
 
+    @expose_as_static
+    def between(
+        self, lower: Expression | float, upper: Expression | float
+    ) -> "BooleanExpression":
+        """Evaluates if the result of this expression is between
+        the lower bound (inclusive) and upper bound (inclusive).
+
+        This is functionally equivalent to performing an `And` operation with
+        `greater_than_or_equal` and `less_than_or_equal`.
+
+        Example:
+            >>> # Check if the 'age' field is between 18 and 65
+            >>> Field.of("age").between(18, 65)
+
+        Args:
+            lower: Lower bound (inclusive) of the range.
+            upper: Upper bound (inclusive) of the range.
+
+        Returns:
+            A new `BooleanExpression` representing the between comparison.
+        """
+        return And(
+            self.greater_than_or_equal(lower),
+            self.less_than_or_equal(upper),
+        )
+
+    @expose_as_static
+    def geo_distance(
+        self, other: Expression | GeoPoint | tuple[float, float]
+    ) -> "FunctionExpression":
+        """Evaluates to the distance in meters between the location in the specified
+        field and the query location.
+
+        Note: This Expression can only be used within a `Search` stage.
+
+        Example:
+            >>> # Calculate distance between the 'location' field and a target GeoPoint
+            >>> Field.of("location").geo_distance(target_point)
+            >>> # Calculate distance between the 'location' field and a (latitude, longitude) tuple
+            >>> Field.of("location").geo_distance((37.7749, -122.4194))
+
+        Args:
+            other: target point used to calculate distance. Can be a GeoPoint, an
+                Expression resolving to a GeoPoint, or a (latitude, longitude) tuple.
+
+        Returns:
+            A new `FunctionExpression` representing the distance.
+        """
+        if isinstance(other, tuple) and len(other) == 2:
+            other = GeoPoint(other[0], other[1])
+
+        return FunctionExpression(
+            "geo_distance", [self, self._cast_to_expr_or_convert_to_constant(other)]
+        )
+
     @expose_as_static
     def equal_any(
         self, array: Array | Sequence[Expression | CONSTANT_TYPE] | Expression
@@ -2927,6 +2982,56 @@ def __init__(self):
         super().__init__("rand", [], use_infix_repr=False)
 
 
+class Score(FunctionExpression):
+    """Evaluates to the search score that reflects the topicality of the document
+    to all of the text predicates (`queryMatch`)
+    in the search query. If `SearchOptions.query` is not set or does not contain
+    any text predicates, then this topicality score will always be `0`.
+
+    Note: This Expression can only be used within a `Search` stage.
+
+    Example:
+        >>> # Sort by search score and retrieve it via add_fields
+        >>> db.pipeline().collection("restaurants").search(
+        ...     query="tacos",
+        ...     sort=Score().descending(),
+        ...     add_fields=[Score().as_("search_score")]
+        ... )
+
+    Returns:
+        A new `Expression` representing the score operation.
+    """
+
+    def __init__(self):
+        super().__init__("score", [], use_infix_repr=False)
+
+
+class DocumentMatches(BooleanExpression):
+    """Creates a boolean expression for a document match query.
+
+    Note: This Expression can only be used within a `Search` stage.
+
+    Example:
+        >>> # Find documents matching the query string
+        >>> db.pipeline().collection("restaurants").search(
+        ...     query=DocumentMatches("pizza OR pasta")
+        ... )
+
+    Args:
+        query: The search query string or expression.
+
+    Returns:
+        A new `BooleanExpression` representing the document match.
+    """
+
+    def __init__(self, query: Expression | str):
+        super().__init__(
+            "document_matches",
+            [Expression._cast_to_expr_or_convert_to_constant(query)],
+            use_infix_repr=False,
+        )
+
+
 class Variable(Expression):
     """
     Creates an expression that retrieves the value of a variable bound via `Pipeline.define`.

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

@@ -30,6 +30,7 @@
     AliasedExpression,
     BooleanExpression,
     CONSTANT_TYPE,
+    DocumentMatches,
     Expression,
     Field,
     Ordering,
@@ -109,6 +110,79 @@ def percentage(value: float):
         return SampleOptions(value, mode=SampleOptions.Mode.PERCENT)
 
 
+class SearchOptions:
+    """Options for configuring the `Search` pipeline stage."""
+
+    def __init__(
+        self,
+        query: str | BooleanExpression,
+        *,
+        limit: Optional[int] = None,
+        retrieval_depth: Optional[int] = None,
+        sort: Optional[Sequence[Ordering] | Ordering] = None,
+        add_fields: Optional[Sequence[Selectable]] = None,
+        offset: Optional[int] = None,
+        language_code: Optional[str] = None,
+    ):
+        """
+        Initializes a SearchOptions instance.
+
+        Args:
+            query (str | BooleanExpression): Specifies the search query that will be used to query and score documents
+                by the search stage. The query can be expressed as an `Expression`, which will be used to score
+                and filter the results. Not all expressions supported by Pipelines are supported in the Search query.
+                The query can also be expressed as a string in the Search DSL.
+            limit (Optional[int]): The maximum number of documents to return from the Search stage.
+            retrieval_depth (Optional[int]): The maximum number of documents for the search stage to score. Documents
+                will be processed in the pre-sort order specified by the search index.
+            sort (Optional[Sequence[Ordering] | Ordering]): Orderings specify how the input documents are sorted.
+            add_fields (Optional[Sequence[Selectable]]): The fields to add to each document, specified as a `Selectable`.
+            offset (Optional[int]): The number of documents to skip.
+            language_code (Optional[str]): The BCP-47 language code of text in the search query, such as "en-US" or "sr-Latn".
+        """
+        self.query = DocumentMatches(query) if isinstance(query, str) else query
+        self.limit = limit
+        self.retrieval_depth = retrieval_depth
+        self.sort = [sort] if isinstance(sort, Ordering) else sort
+        self.add_fields = add_fields
+        self.offset = offset
+        self.language_code = language_code
+
+    def __repr__(self):
+        args = [f"query={self.query!r}"]
+        if self.limit is not None:
+            args.append(f"limit={self.limit}")
+        if self.retrieval_depth is not None:
+            args.append(f"retrieval_depth={self.retrieval_depth}")
+        if self.sort is not None:
+            args.append(f"sort={self.sort}")
+        if self.add_fields is not None:
+            args.append(f"add_fields={self.add_fields}")
+        if self.offset is not None:
+            args.append(f"offset={self.offset}")
+        if self.language_code is not None:
+            args.append(f"language_code={self.language_code!r}")
+        return f"{self.__class__.__name__}({', '.join(args)})"
+
+    def _to_dict(self) -> dict[str, Value]:
+        options = {"query": self.query._to_pb()}
+        if self.limit is not None:
+            options["limit"] = Value(integer_value=self.limit)
+        if self.retrieval_depth is not None:
+            options["retrieval_depth"] = Value(integer_value=self.retrieval_depth)
+        if self.sort is not None:
+            options["sort"] = Value(
+                array_value={"values": [s._to_pb() for s in self.sort]}
+            )
+        if self.add_fields is not None:
+            options["add_fields"] = Selectable._to_value(self.add_fields)
+        if self.offset is not None:
+            options["offset"] = Value(integer_value=self.offset)
+        if self.language_code is not None:
+            options["language_code"] = Value(string_value=self.language_code)
+        return options
+
+
 class UnnestOptions:
     """Options for configuring the `Unnest` pipeline stage.
 
@@ -423,6 +497,24 @@ def _pb_args(self):
             ]
 
 
+class Search(Stage):
+    """Search stage."""
+
+    def __init__(self, query_or_options: str | BooleanExpression | SearchOptions):
+        super().__init__("search")
+        if isinstance(query_or_options, SearchOptions):
+            options = query_or_options
+        else:
+            options = SearchOptions(query=query_or_options)
+        self.options = options
+
+    def _pb_args(self) -> list[Value]:
+        return []
+
+    def _pb_options(self) -> dict[str, Value]:
+        return self.options._to_dict()
+
+
 class Select(Stage):
     """Selects or creates a set of fields."""
 

packages/google-cloud-firestore/tests/system/pipeline_e2e/data.yaml

@@ -148,8 +148,13 @@ data:
   cities:
     city1:
       name: "San Francisco"
+      location: GEOPOINT(37.7749,-122.4194)
     city2:
       name: "New York"
+      location: GEOPOINT(40.7128,-74.0060)
+    city3:
+      name: "Saskatoon"
+      location: GEOPOINT(52.1579,-106.6702)
   "cities/city1/landmarks":
     lm1:
       name: "Golden Gate Bridge"
@@ -167,4 +172,57 @@ data:
       rating: 5
     rev2:
       author: "Bob"
-      rating: 4
+      rating: 4
+  "cities/city3/landmarks":
+    lm4:
+      name: "Western Development Museum"
+      type: "Museum"
+  restaurants:
+    sunnySideUp:
+      name: "The Sunny Side Up"
+      description: "A cozy neighborhood diner serving classic breakfast favorites all day long, from fluffy pancakes to savory omelets."
+      location: GEOPOINT(39.7541,-105.0002)
+      menu: "<h3>Breakfast Classics</h3><ul><li>Denver Omelet - $12</li><li>Buttermilk Pancakes - $10</li><li>Steak and Eggs - $16</li></ul><h3>Sides</h3><ul><li>Hash Browns - $4</li><li>Thick-cut Bacon - $5</li><li>Drip Coffee - $2</li></ul>"
+      average_price_per_person: 15
+    goldenWaffle:
+      name: "The Golden Waffle"
+      description: "Specializing exclusively in Belgian-style waffles. Open daily from 6:00 AM to 11:00 AM."
+      location: GEOPOINT(39.7183,-104.9621)
+      menu: "<h3>Signature Waffles</h3><ul><li>Strawberry Delight - $11</li><li>Chicken and Waffles - $14</li><li>Chocolate Chip Crunch - $10</li></ul><h3>Drinks</h3><ul><li>Fresh OJ - $4</li><li>Artisan Coffee - $3</li></ul>"
+      average_price_per_person: 13
+    lotusBlossomThai:
+      name: "Lotus Blossom Thai"
+      description: "Authentic Thai cuisine featuring hand-crushed spices and traditional family recipes from the Chiang Mai region."
+      location: GEOPOINT(39.7315,-104.9847)
+      menu: "<h3>Appetizers</h3><ul><li>Spring Rolls - $7</li><li>Chicken Satay - $9</li></ul><h3>Main Course</h3><ul><li>Pad Thai - $15</li><li>Green Curry - $16</li><li>Drunken Noodles - $15</li></ul>"
+      average_price_per_person: 22
+    mileHighCatch:
+      name: "Mile High Catch"
+      description: "Freshly sourced seafood offering a wide variety of Pacific fish and Atlantic shellfish in an upscale atmosphere."
+      location: GEOPOINT(39.7401,-104.9903)
+      menu: "<h3>From the Raw Bar</h3><ul><li>Oysters (Half Dozen) - $18</li><li>Lobster Cocktail - $22</li></ul><h3>Entrees</h3><ul><li>Pan-Seared Salmon - $28</li><li>King Crab Legs - $45</li><li>Fish and Chips - $19</li></ul>"
+      average_price_per_person: 45
+    peakBurgers:
+      name: "Peak Burgers"
+      description: "Casual burger joint focused on locally sourced Colorado beef and hand-cut fries."
+      location: GEOPOINT(39.7622,-105.0125)
+      menu: "<h3>Burgers</h3><ul><li>The Peak Double - $12</li><li>Bison Burger - $15</li><li>Veggie Stack - $11</li></ul><h3>Sides</h3><ul><li>Truffle Fries - $6</li><li>Onion Rings - $5</li></ul>"
+      average_price_per_person: 18
+    solTacos:
+      name: "El Sol Tacos"
+      description: "A vibrant street-side taco stand serving up quick, delicious, and traditional Mexican street food."
+      location: GEOPOINT(39.6952,-105.0274)
+      menu: "<h3>Tacos ($3.50 each)</h3><ul><li>Al Pastor</li><li>Carne Asada</li><li>Pollo Asado</li><li>Nopales (Cactus)</li></ul><h3>Beverages</h3><ul><li>Horchata - $4</li><li>Mexican Coke - $3</li></ul>"
+      average_price_per_person: 12
+    eastsideTacos:
+      name: "Eastside Cantina"
+      description: "Authentic street tacos and hand-shaken margaritas on the vibrant east side of the city."
+      location: GEOPOINT(39.735,-104.885)
+      menu: "<h3>Tacos</h3><ul><li>Carnitas Tacos - $4</li><li>Barbacoa Tacos - $4.50</li><li>Shrimp Tacos - $5</li></ul><h3>Drinks</h3><ul><li>House Margarita - $9</li><li>Jarritos - $3</li></ul>"
+      average_price_per_person: 18
+    eastsideChicken:
+      name: "Eastside Chicken"
+      description: "Fried chicken to go - next to Eastside Cantina."
+      location: GEOPOINT(39.735,-104.885)
+      menu: "<h3>Fried Chicken</h3><ul><li>Drumstick - $4</li><li>Wings - $1</li><li>Sandwich - $9</li></ul><h3>Drinks</h3><ul><li>House Margarita - $9</li><li>Jarritos - $3</li></ul>"
+      average_price_per_person: 12