added array_get

Author: daniel-sanche

google/cloud/firestore_v1/pipeline_expressions.py

@@ -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/unit/v1/test_pipeline_expressions.py

@@ -706,6 +706,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")