initial commit

Author: daniel-sanche

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

@@ -344,6 +344,20 @@ def sort(self, *orders: stages.Ordering) -> "_BasePipeline":
         """
         return self._append(stages.Sort(*orders))
 
+    def search(self, options: stages.SearchOptions) -> "_BasePipeline":
+        """
+        Adds a search stage to the pipeline.
+
+        This stage filters documents based on the provided query expression.
+
+        Args:
+            options: A SearchOptions instance configuring the search.
+
+        Returns:
+            A new Pipeline object with this stage appended to the stage list
+        """
+        return self._append(stages.Search(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

@@ -609,6 +609,50 @@ def less_than_or_equal(
             [self, self._cast_to_expr_or_convert_to_constant(other)],
         )
 
+    @expose_as_static
+    def between(
+        self, lower: Expression | CONSTANT_TYPE, upper: Expression | CONSTANT_TYPE
+    ) -> "BooleanExpression":
+        """Creates an expression that checks if this expression is between two values.
+
+        Example:
+            >>> # Check if the 'age' field is between 18 and 65
+            >>> Field.of("age").between(18, 65)
+
+        Args:
+            lower: The lower bound expression or constant value.
+            upper: The upper bound expression or constant value.
+
+        Returns:
+            A new `BooleanExpression` representing the between comparison.
+        """
+        return BooleanExpression(
+            "between",
+            [
+                self,
+                self._cast_to_expr_or_convert_to_constant(lower),
+                self._cast_to_expr_or_convert_to_constant(upper),
+            ],
+        )
+
+    @expose_as_static
+    def geo_distance(self, other: Expression | CONSTANT_TYPE) -> "FunctionExpression":
+        """Creates an expression that calculates the distance between two geographical points.
+
+        Example:
+            >>> # Calculate distance between the 'location' field and a target GeoPoint
+            >>> Field.of("location").geo_distance(target_point)
+
+        Args:
+            other: The other point expression or constant value.
+
+        Returns:
+            A new `FunctionExpression` representing the distance.
+        """
+        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
@@ -2634,3 +2678,22 @@ class Rand(FunctionExpression):
 
     def __init__(self):
         super().__init__("rand", [], use_infix_repr=False)
+
+
+def document_matches(query: Expression | str) -> BooleanExpression:
+    """Creates a boolean expression for a document match query.
+    
+    Example:
+        >>> # Find documents matching the query string
+        >>> document_matches("search query")
+
+    Args:
+        query: The search query string or expression.
+
+    Returns:
+        A new `BooleanExpression` representing the document match.
+    """
+    return BooleanExpression(
+        "document_matches", [Expression._cast_to_expr_or_convert_to_constant(query)]
+    )
+

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

@@ -109,6 +109,85 @@ def percentage(value: float):
         return SampleOptions(value, mode=SampleOptions.Mode.PERCENT)
 
 
+class QueryEnhancement(Enum):
+    """Define the query expansion behavior used by full-text search expressions."""
+    DISABLED = "disabled"
+    REQUIRED = "required"
+    PREFERRED = "preferred"
+
+
+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,
+        select: Optional[Sequence[Selectable | str]] = None,
+        offset: Optional[int] = None,
+        query_enhancement: Optional[str | QueryEnhancement] = 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`.
+            select (Optional[Sequence[Selectable | str]]): The fields to keep or add to each document,
+                specified as an array of `Selectable` or strings.
+            offset (Optional[int]): The number of documents to skip.
+            query_enhancement (Optional[str | QueryEnhancement]): Define the query expansion behavior used by full-text search expressions
+                in this search stage.
+            language_code (Optional[str]): The BCP-47 language code of text in the search query, such as "en-US" or "sr-Latn".
+        """
+        if isinstance(query, str):
+            from google.cloud.firestore_v1.pipeline_expressions import document_matches
+            self.query = document_matches(query)
+        else:
+            self.query = query
+        self.limit = limit
+        self.retrieval_depth = retrieval_depth
+        self.sort = [sort] if isinstance(sort, Ordering) else sort
+        self.add_fields = add_fields
+        self.select = select
+        self.offset = offset
+        self.query_enhancement = (
+            QueryEnhancement(query_enhancement.lower()) if isinstance(query_enhancement, str) else query_enhancement
+        )
+        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.select is not None:
+            args.append(f"select={self.select}")
+        if self.offset is not None:
+            args.append(f"offset={self.offset}")
+        if self.query_enhancement is not None:
+            args.append(f"query_enhancement={self.query_enhancement!r}")
+        if self.language_code is not None:
+            args.append(f"language_code={self.language_code!r}")
+        return f"{self.__class__.__name__}({', '.join(args)})"
+
+
 class UnnestOptions:
     """Options for configuring the `Unnest` pipeline stage.
 
@@ -423,6 +502,41 @@ def _pb_args(self):
             ]
 
 
+class Search(Stage):
+    """Search stage."""
+
+    def __init__(self, options: SearchOptions):
+        super().__init__("search")
+        self.options = options
+
+    def _pb_args(self) -> list[Value]:
+        return []
+
+    def _pb_options(self) -> dict[str, Value]:
+        options = {}
+        if self.options.query is not None:
+            options["query"] = self.options.query._to_pb()
+        if self.options.limit is not None:
+            options["limit"] = Value(integer_value=self.options.limit)
+        if self.options.retrieval_depth is not None:
+            options["retrieval_depth"] = Value(integer_value=self.options.retrieval_depth)
+        if self.options.sort is not None:
+            options["sort"] = Value(array_value={"values": [s._to_pb() for s in self.options.sort]})
+        if self.options.add_fields is not None:
+            from google.cloud.firestore_v1.pipeline_expressions import Selectable
+            options["add_fields"] = Selectable._to_value(self.options.add_fields)
+        if self.options.select is not None:
+            from google.cloud.firestore_v1.pipeline_expressions import Selectable
+            options["select"] = Selectable._to_value(self.options.select)
+        if self.options.offset is not None:
+            options["offset"] = Value(integer_value=self.options.offset)
+        if self.options.query_enhancement is not None:
+            options["query_enhancement"] = Value(string_value=self.options.query_enhancement.value)
+        if self.options.language_code is not None:
+            options["language_code"] = Value(string_value=self.options.language_code)
+        return options
+
+
 class Select(Stage):
     """Selects or creates a set of fields."""
 

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

@@ -144,4 +144,14 @@ data:
     doc_with_nan:
       value: "NaN"
     doc_with_null:
-      value: null
+      value: null
+  geopoints:
+    loc1:
+      name: SF
+      location: GEOPOINT(37.7749,-122.4194)
+    loc2:
+      name: LA
+      location: GEOPOINT(34.0522,-118.2437)
+    loc3:
+      name: NY
+      location: GEOPOINT(40.7128,-74.0060)

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

@@ -0,0 +1,24 @@
+tests:
+  - description: search_stage_basic
+    pipeline:
+      - Collection: books
+      - Search:
+          - SearchOptions:
+              query: "technology"
+              limit: 2
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args: []
+          name: search
+          options:
+            limit:
+              integerValue: '2'
+            query:
+              functionValue:
+                args:
+                - stringValue: "technology"
+                name: document_matches

packages/google-cloud-firestore/tests/system/test_pipeline_acceptance.py

@@ -33,6 +33,7 @@
 from google.cloud.firestore_v1 import pipeline_expressions as expr
 from google.cloud.firestore_v1 import pipeline_stages as stages
 from google.cloud.firestore_v1.vector import Vector
+from google.cloud.firestore_v1 import GeoPoint
 
 FIRESTORE_PROJECT = os.environ.get("GCLOUD_PROJECT")
 
@@ -343,12 +344,15 @@ def _parse_yaml_types(data):
         else:
             return [_parse_yaml_types(value) for value in data]
     # detect timestamps
-    if isinstance(data, str) and ":" in data:
+    if isinstance(data, str) and ":" in data and not data.startswith("GEOPOINT("):
         try:
             parsed_datetime = datetime.datetime.fromisoformat(data)
             return parsed_datetime
         except ValueError:
             pass
+    if isinstance(data, str) and data.startswith("GEOPOINT("):
+        parts = data[9:-1].split(",")
+        return GeoPoint(float(parts[0]), float(parts[1]))
     if data == "NaN":
         return float("NaN")
     return data

packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py

@@ -790,6 +790,34 @@ def test_equal(self):
         infix_instance = arg1.equal(arg2)
         assert infix_instance == instance
 
+    def test_between(self):
+        arg1 = self._make_arg("Left")
+        arg2 = self._make_arg("Lower")
+        arg3 = self._make_arg("Upper")
+        instance = Expression.between(arg1, arg2, arg3)
+        assert instance.name == "between"
+        assert instance.params == [arg1, arg2, arg3]
+        assert repr(instance) == "Left.between(Lower, Upper)"
+        infix_instance = arg1.between(arg2, arg3)
+        assert infix_instance == instance
+
+    def test_geo_distance(self):
+        arg1 = self._make_arg("Left")
+        arg2 = self._make_arg("Right")
+        instance = Expression.geo_distance(arg1, arg2)
+        assert instance.name == "geo_distance"
+        assert instance.params == [arg1, arg2]
+        assert repr(instance) == "Left.geo_distance(Right)"
+        infix_instance = arg1.geo_distance(arg2)
+        assert infix_instance == instance
+
+    def test_document_matches(self):
+        arg1 = self._make_arg("Query")
+        instance = expr.document_matches(arg1)
+        assert instance.name == "document_matches"
+        assert instance.params == [arg1]
+        assert repr(instance) == "document_matches(Query)"
+
     def test_greater_than_or_equal(self):
         arg1 = self._make_arg("Left")
         arg2 = self._make_arg("Right")

packages/google-cloud-firestore/tests/unit/v1/test_pipeline_search.py

@@ -0,0 +1,53 @@
+from google.cloud.firestore_v1.pipeline_stages import SearchOptions, Search, QueryEnhancement
+from google.cloud.firestore_v1.pipeline_expressions import Field, Ordering, document_matches
+
+def test_search_options():
+    options = SearchOptions(
+        query="test query",
+        limit=10,
+        retrieval_depth=2,
+        sort=Ordering("score", Ordering.Direction.DESCENDING),
+        add_fields=[Field("extra")],
+        select=[Field("name")],
+        offset=5,
+        query_enhancement="disabled",
+        language_code="en",
+    )
+    assert options.limit == 10
+    assert options.retrieval_depth == 2
+    assert len(options.sort) == 1
+    assert options.offset == 5
+    assert options.query_enhancement == QueryEnhancement.DISABLED
+    assert options.language_code == "en"
+    
+    # Check proto generation
+    stage = Search(options)
+    pb_opts = stage._pb_options()
+    
+    assert pb_opts["limit"].integer_value == 10
+    assert pb_opts["retrieval_depth"].integer_value == 2
+    assert len(pb_opts["sort"].array_value.values) == 1
+    assert pb_opts["offset"].integer_value == 5
+    assert pb_opts["query_enhancement"].string_value == "disabled"
+    assert pb_opts["language_code"].string_value == "en"
+
+def test_search_options_bool_expr():
+    options = SearchOptions(query=document_matches("query string"))
+    stage = Search(options)
+    pb_opts = stage._pb_options()
+    assert "query" in pb_opts
+
+def test_between():
+    expr = Field("age").between(18, 65)
+    assert expr.name == "between"
+    assert len(expr.params) == 3
+
+def test_geo_distance():
+    expr = Field("location").geo_distance("other")
+    assert expr.name == "geo_distance"
+    assert len(expr.params) == 2
+
+def test_document_matches():
+    expr = document_matches("search query")
+    assert expr.name == "document_matches"
+    assert len(expr.params) == 1