Merge branch 'pipeline-preview' into explain_stats

Author: daniel-sanche

google/cloud/firestore_v1/pipeline_expressions.py

@@ -396,7 +396,7 @@ def sqrt(self) -> "Expression":
         return Function("sqrt", [self])
 
     @expose_as_static
-    def logical_maximum(self, other: Expression | CONSTANT_TYPE) -> "Expression":
+    def logical_maximum(self, *others: Expression | CONSTANT_TYPE) -> "Expression":
         """Creates an expression that returns the larger value between this expression
         and another expression or constant, based on Firestore's value type ordering.
 
@@ -406,23 +406,23 @@ def logical_maximum(self, other: Expression | CONSTANT_TYPE) -> "Expression":
         Example:
             >>> # Returns the larger value between the 'discount' field and the 'cap' field.
             >>> Field.of("discount").logical_maximum(Field.of("cap"))
-            >>> # Returns the larger value between the 'value' field and 10.
-            >>> Field.of("value").logical_maximum(10)
+            >>> # Returns the larger value between the 'value' field and some ints
+            >>> Field.of("value").logical_maximum(10, 20, 30)
 
         Args:
-            other: The other expression or constant value to compare with.
+            others: The other expression or constant values to compare with.
 
         Returns:
             A new `Expression` representing the logical maximum operation.
         """
         return Function(
             "maximum",
-            [self, self._cast_to_expr_or_convert_to_constant(other)],
+            [self] + [self._cast_to_expr_or_convert_to_constant(o) for o in others],
             infix_name_override="logical_maximum",
         )
 
     @expose_as_static
-    def logical_minimum(self, other: Expression | CONSTANT_TYPE) -> "Expression":
+    def logical_minimum(self, *others: Expression | CONSTANT_TYPE) -> "Expression":
         """Creates an expression that returns the smaller value between this expression
         and another expression or constant, based on Firestore's value type ordering.
 
@@ -432,18 +432,18 @@ def logical_minimum(self, other: Expression | CONSTANT_TYPE) -> "Expression":
         Example:
             >>> # Returns the smaller value between the 'discount' field and the 'floor' field.
             >>> Field.of("discount").logical_minimum(Field.of("floor"))
-            >>> # Returns the smaller value between the 'value' field and 10.
-            >>> Field.of("value").logical_minimum(10)
+            >>> # Returns the smaller value between the 'value' field and some ints
+            >>> Field.of("value").logical_minimum(10, 20, 30)
 
         Args:
-            other: The other expression or constant value to compare with.
+            others: The other expression or constant values to compare with.
 
         Returns:
             A new `Expression` representing the logical minimum operation.
         """
         return Function(
             "minimum",
-            [self, self._cast_to_expr_or_convert_to_constant(other)],
+            [self] + [self._cast_to_expr_or_convert_to_constant(o) for o in others],
             infix_name_override="logical_minimum",
         )
 
@@ -629,6 +629,25 @@ def not_equal_any(
             ],
         )
 
+    @expose_as_static
+    def array_get(self, offset: Expression | int) -> "Function":
+        """
+        Creates an expression that indexes into an array from the beginning or end and returns the
+        element. A negative offset starts from the end.
+
+        Example:
+            >>> Array([1,2,3]).array_get(0)
+
+        Args:
+            offset: the index of the element to return
+
+        Returns:
+            A new `Expression` representing the `array_get` operation.
+        """
+        return Function(
+            "array_get", [self, self._cast_to_expr_or_convert_to_constant(offset)]
+        )
+
     @expose_as_static
     def array_contains(
         self, element: Expression | CONSTANT_TYPE

tests/system/pipeline_e2e/array.yaml

@@ -386,3 +386,79 @@ tests:
                         name: array
                     name: array_concat
           name: select
+  - description: testArrayGet
+    pipeline:
+      - Collection: books
+      - Where:
+          - Function.equal:
+              - Field: title
+              - Constant: "The Hitchhiker's Guide to the Galaxy"
+      - Select:
+        - AliasedExpression:
+            - Function.array_get:
+                - Field: tags
+                - Constant: 0
+            - "firstTag"
+    assert_results:
+      - firstTag: "comedy"
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args:
+          - functionValue:
+              args:
+              - fieldReferenceValue: title
+              - stringValue: "The Hitchhiker's Guide to the Galaxy"
+              name: equal
+          name: where
+        - args:
+          - mapValue:
+              fields:
+                firstTag:
+                  functionValue:
+                    args:
+                    - fieldReferenceValue: tags
+                    - integerValue: '0'
+                    name: array_get
+          name: select
+  - description: testArrayGet_NegativeOffset
+    pipeline:
+      - Collection: books
+      - Where:
+          - Function.equal:
+              - Field: title
+              - Constant: "The Hitchhiker's Guide to the Galaxy"
+      - Select:
+        - AliasedExpression:
+            - Function.array_get:
+                - Field: tags
+                - Constant: -1
+            - "lastTag"
+    assert_results:
+      - lastTag: "adventure"
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args:
+          - functionValue:
+              args:
+              - fieldReferenceValue: title
+              - stringValue: "The Hitchhiker's Guide to the Galaxy"
+              name: equal
+          name: where
+        - args:
+          - mapValue:
+              fields:
+                lastTag:
+                  functionValue:
+                    args:
+                    - fieldReferenceValue: tags
+                    - integerValue: '-1'
+                    name: array_get
+          name: select

tests/system/pipeline_e2e/logical.yaml

@@ -464,6 +464,64 @@ tests:
                     - doubleValue: 4.5
                     name: maximum
           name: select
+  - description: testLogicalMinMaxWithMultipleInputs
+    pipeline:
+      - Collection: books
+      - Where:
+          - Function.equal:
+              - Field: author
+              - Constant: Douglas Adams
+      - Select:
+          - AliasedExpression:
+              - Function.logical_maximum:
+                  - Field: rating
+                  - Constant: 4.5
+                  - Constant: 3.0
+                  - Constant: 5.0
+              - "max_rating"
+          - AliasedExpression:
+              - Function.logical_minimum:
+                  - Field: published
+                  - Constant: 1900
+                  - Constant: 2000
+                  - Constant: 1984
+              - "min_published"
+    assert_results:
+      - max_rating: 5.0
+        min_published: 1900
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args:
+          - functionValue:
+              args:
+              - fieldReferenceValue: author
+              - stringValue: Douglas Adams
+              name: equal
+          name: where
+        - args:
+          - mapValue:
+              fields:
+                min_published:
+                  functionValue:
+                    args:
+                    - fieldReferenceValue: published
+                    - integerValue: '1900'
+                    - integerValue: '2000'
+                    - integerValue: '1984'
+                    name: minimum
+                max_rating:
+                  functionValue:
+                    args:
+                    - fieldReferenceValue: rating
+                    - doubleValue: 4.5
+                    - doubleValue: 3.0
+                    - doubleValue: 5.0
+                    name: maximum
+          name: select
   - description: testGreaterThanOrEqual
     pipeline:
       - Collection: books

tests/unit/v1/test_pipeline_expressions.py

@@ -717,6 +717,16 @@ def test_or(self):
         assert instance.params == [arg1, arg2]
         assert repr(instance) == "Or(Arg1, Arg2)"
 
+    def test_array_get(self):
+        arg1 = self._make_arg("ArrayField")
+        arg2 = self._make_arg("Offset")
+        instance = Expression.array_get(arg1, arg2)
+        assert instance.name == "array_get"
+        assert instance.params == [arg1, arg2]
+        assert repr(instance) == "ArrayField.array_get(Offset)"
+        infix_istance = arg1.array_get(arg2)
+        assert infix_istance == instance
+
     def test_array_contains(self):
         arg1 = self._make_arg("ArrayField")
         arg2 = self._make_arg("Element")
@@ -989,23 +999,25 @@ def test_divide(self):
         assert infix_instance == instance
 
     def test_logical_maximum(self):
-        arg1 = self._make_arg("Left")
-        arg2 = self._make_arg("Right")
-        instance = Expression.logical_maximum(arg1, arg2)
+        arg1 = self._make_arg("A1")
+        arg2 = self._make_arg("A2")
+        arg3 = self._make_arg("A3")
+        instance = Expression.logical_maximum(arg1, arg2, arg3)
         assert instance.name == "maximum"
-        assert instance.params == [arg1, arg2]
-        assert repr(instance) == "Left.logical_maximum(Right)"
-        infix_instance = arg1.logical_maximum(arg2)
+        assert instance.params == [arg1, arg2, arg3]
+        assert repr(instance) == "A1.logical_maximum(A2, A3)"
+        infix_instance = arg1.logical_maximum(arg2, arg3)
         assert infix_instance == instance
 
     def test_logical_minimum(self):
-        arg1 = self._make_arg("Left")
-        arg2 = self._make_arg("Right")
-        instance = Expression.logical_minimum(arg1, arg2)
+        arg1 = self._make_arg("A1")
+        arg2 = self._make_arg("A2")
+        arg3 = self._make_arg("A3")
+        instance = Expression.logical_minimum(arg1, arg2, arg3)
         assert instance.name == "minimum"
-        assert instance.params == [arg1, arg2]
-        assert repr(instance) == "Left.logical_minimum(Right)"
-        infix_instance = arg1.logical_minimum(arg2)
+        assert instance.params == [arg1, arg2, arg3]
+        assert repr(instance) == "A1.logical_minimum(A2, A3)"
+        infix_instance = arg1.logical_minimum(arg2, arg3)
         assert infix_instance == instance
 
     def test_to_lower(self):