feat(firestore): support array functions

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

@@ -1490,6 +1490,250 @@ def last(self) -> "Expression":
         """
         return AggregateFunction("last", [self])
 
+    @expose_as_static
+    def array_first(self) -> "Expression":
+        """Creates an expression that returns the first element of an array.
+
+        Example:
+            >>> # Select the first element of array 'colors'
+            >>> Field.of("colors").array_first()
+
+        Returns:
+            A new `Expression` representing the first element of the array.
+        """
+        return FunctionExpression("array_first", [self])
+
+    @expose_as_static
+    def array_last(self) -> "Expression":
+        """Creates an expression that returns the last element of an array.
+
+        Example:
+            >>> # Select the last element of array 'colors'
+            >>> Field.of("colors").array_last()
+
+        Returns:
+            A new `Expression` representing the last element of the array.
+        """
+        return FunctionExpression("array_last", [self])
+
+    @expose_as_static
+    def array_first_n(self, n: int | "Expression") -> "Expression":
+        """Creates an expression that returns the first `n` elements of an array.
+
+        Example:
+            >>> # Select the first 2 elements of array 'colors'
+            >>> Field.of("colors").array_first_n(2)
+
+        Returns:
+            A new `Expression` representing the first `n` elements of the array.
+        """
+        return FunctionExpression(
+            "array_first_n", [self, self._cast_to_expr_or_convert_to_constant(n)]
+        )
+
+    @expose_as_static
+    def array_last_n(self, n: int | "Expression") -> "Expression":
+        """Creates an expression that returns the last `n` elements of an array.
+
+        Example:
+            >>> # Select the last 2 elements of array 'colors'
+            >>> Field.of("colors").array_last_n(2)
+
+        Returns:
+            A new `Expression` representing the last `n` elements of the array.
+        """
+        return FunctionExpression(
+            "array_last_n", [self, self._cast_to_expr_or_convert_to_constant(n)]
+        )
+
+    @expose_as_static
+    def array_maximum(self) -> "Expression":
+        """Creates an expression that returns the maximum element of an array.
+
+        Example:
+            >>> # Select the maximum element of array 'scores'
+            >>> Field.of("scores").array_maximum()
+
+        Returns:
+            A new `Expression` representing the maximum element of the array.
+        """
+        return FunctionExpression("maximum", [self])
+
+    @expose_as_static
+    def array_minimum(self) -> "Expression":
+        """Creates an expression that returns the minimum element of an array.
+
+        Example:
+            >>> # Select the minimum element of array 'scores'
+            >>> Field.of("scores").array_minimum()
+
+        Returns:
+            A new `Expression` representing the minimum element of the array.
+        """
+        return FunctionExpression("minimum", [self])
+
+    @expose_as_static
+    def array_maximum_n(self, n: int | "Expression") -> "Expression":
+        """Creates an expression that returns the maximum `n` elements of an array.
+
+        Example:
+            >>> # Select the maximum 2 elements of array 'scores'
+            >>> Field.of("scores").array_maximum_n(2)
+
+        Note:
+            Returns the n largest non-null elements in the array, in descending
+            order. This does not use a stable sort, meaning the order of equivalent
+            elements is undefined.
+
+        Returns:
+            A new `Expression` representing the maximum `n` elements of the array.
+        """
+        return FunctionExpression(
+            "maximum_n", [self, self._cast_to_expr_or_convert_to_constant(n)]
+        )
+
+    @expose_as_static
+    def array_minimum_n(self, n: int | "Expression") -> "Expression":
+        """Creates an expression that returns the minimum `n` elements of an array.
+
+        Example:
+            >>> # Select the minimum 2 elements of array 'scores'
+            >>> Field.of("scores").array_minimum_n(2)
+
+        Note:
+            Returns the n smallest non-null elements in the array, in ascending
+            order. This does not use a stable sort, meaning the order of equivalent
+            elements is undefined.
+
+        Returns:
+            A new `Expression` representing the minimum `n` elements of the array.
+        """
+        return FunctionExpression(
+            "minimum_n", [self, self._cast_to_expr_or_convert_to_constant(n)]
+        )
+
+    @expose_as_static
+    def array_slice(
+        self, offset: int | "Expression", length: int | "Expression" | None = None
+    ) -> "Expression":
+        """Ccreates an expression that returns a slice of an array starting from the specified
+        offset with a given length.
+
+        Example:
+            >>> # Slice array 'scores' starting at index 1 with length 2
+            >>> Field.of("scores").array_slice(1, 2)
+
+        Args:
+            offset: the 0-based index of the first element to include.
+            length: The number of elements to include in the slice. If omitted, slices to the end.
+
+        Returns:
+            A new `Expression` representing the slice of the array.
+        """
+        args = [self, self._cast_to_expr_or_convert_to_constant(offset)]
+        if length is not None:
+            args.append(self._cast_to_expr_or_convert_to_constant(length))
+        return FunctionExpression("array_slice", args)
+
+    @expose_as_static
+    def array_index_of(self, search: "Expression" | CONSTANT_TYPE) -> "Expression":
+        """Creates an expression that returns the index of a value in an array.
+
+        Returns -1 if the value is not found.
+
+        Example:
+            >>> # Get the index of "comedy" in the 'tags' array
+            >>> Field.of("tags").array_index_of("comedy")
+
+        Args:
+            search: The element (expression or constant) to find the index of.
+
+        Returns:
+            A new `Expression` representing the 'array_index_of' value.
+        """
+        return FunctionExpression(
+            "array_index_of",
+            [
+                self,
+                self._cast_to_expr_or_convert_to_constant(search),
+                self._cast_to_expr_or_convert_to_constant("first"),
+            ],
+        )
+
+    @expose_as_static
+    def array_index_of_all(self, search: "Expression" | CONSTANT_TYPE) -> "Expression":
+        """Creates an expression that returns all indices of a value in an array.
+        Returns an empty array if the value is not found.
+
+        Example:
+            >>> # Get all indices of "comedy" in the 'tags' array
+            >>> Field.of("tags").array_index_of_all("comedy")
+
+        Args:
+            search: The element (expression or constant) to find the indices of.
+
+        Returns:
+            A new `Expression` representing the 'array_index_of_all' value.
+        """
+        return FunctionExpression(
+            "array_index_of_all",
+            [self, self._cast_to_expr_or_convert_to_constant(search)],
+        )
+
+    @expose_as_static
+    def array_transform(
+        self, element_alias: str, body: "Expression", index_alias: str | None = None
+    ) -> "Expression":
+        """Creates an expression that transforms elements of an array.
+
+        Example:
+            >>> # Transform each element by adding 1.
+            >>> Field.of("nums").array_transform("e", Field.of("e").add(1))
+            >>>
+            >>> # Transform each element by adding its index to it.
+            >>> Field.of("nums").array_transform("e", Field.of("e").add(Field.of("i")), index_alias="i")
+
+        Args:
+            element_alias: The variable name to use for the current element within the body expression.
+            body: The expression to apply to each element.
+            index_alias: The variable name to use for the current index within the body expression.
+
+        Returns:
+            A new `Expression` applying the transformation to the array elements.
+        """
+        args = [
+            self,
+            self._cast_to_expr_or_convert_to_constant(element_alias),
+        ]
+        if index_alias is not None:
+            args.append(self._cast_to_expr_or_convert_to_constant(index_alias))
+
+        args.append(self._cast_to_expr_or_convert_to_constant(body))
+
+        return FunctionExpression("array_transform", args)
+
+    @expose_as_static
+    def array_filter(self, element_alias: str, body: "Expression") -> "Expression":
+        """
+        Takes an array, evaluates a boolean expression on each element, and returns a new
+        array containing only the elements for which the expression evaluates to True.
+
+        Args:
+           element_alias: Element variable name.
+           body: Boolean expression applied to each element.
+
+        Returns:
+            Expression: The created FunctionExpression AST node.
+        """
+        return FunctionExpression(
+            "array_filter",
+            [
+                self,
+                self._cast_to_expr_or_convert_to_constant(element_alias),
+                self._cast_to_expr_or_convert_to_constant(body),
+            ],
+        )
+
     @expose_as_static
     def unix_micros_to_timestamp(self) -> "Expression":
         """Creates an expression that converts a number of microseconds since the epoch (1970-01-01