feat(firestore): support object functions (#16132)

* map_set()
* map_keys()
* map_values()
* map_entries()

Author: Linchin

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

@@ -1285,6 +1285,33 @@ def map_get(self, key: str | Constant[str]) -> "Expression":
             "map_get", [self, self._cast_to_expr_or_convert_to_constant(key)]
         )
 
+    def map_set(self, key: str | Constant[str], value: Any) -> "Expression":
+        """Creates an expression that returns a new map with the specified entries added or
+        updated.
+
+        Note:
+            `map_set` only performs shallow updates to the map. Setting a value to `None`
+            will retain the key with a `None` value. To remove a key entirely, use
+            `map_remove`.
+
+        Example:
+            >>> Map({"city": "London"}).map_set("city", "New York")
+            >>> Field.of("address").map_set("city", "Seattle")
+
+        Args:
+            key: The key to set in the map.
+            value: The value to associate with the key.
+
+        Returns:
+            A new `Expression` representing the map_set operation.
+        """
+        args = [
+            self,
+            self._cast_to_expr_or_convert_to_constant(key),
+            self._cast_to_expr_or_convert_to_constant(value),
+        ]
+        return FunctionExpression("map_set", args)
+
     @expose_as_static
     def map_remove(self, key: str | Constant[str]) -> "Expression":
         """Remove a key from a the map produced by evaluating this expression.
@@ -1328,6 +1355,58 @@ def map_merge(
         )
 
     @expose_as_static
+    def map_keys(self) -> "Expression":
+        """Creates an expression that returns the keys of a map.
+
+        Note:
+            While the backend generally preserves insertion order, relying on the
+            order of the output array is not guaranteed and should be avoided.
+
+        Example:
+            >>> Map({"city": "London", "country": "UK"}).map_keys()
+            >>> Field.of("address").map_keys()
+
+        Returns:
+            A new `Expression` representing the keys of the map.
+        """
+        return FunctionExpression("map_keys", [self])
+
+    @expose_as_static
+    def map_values(self) -> "Expression":
+        """Creates an expression that returns the values of a map.
+
+        Note:
+            While the backend generally preserves insertion order, relying on the
+            order of the output array is not guaranteed and should be avoided.
+
+        Example:
+            >>> Map({"city": "London", "country": "UK"}).map_values()
+            >>> Field.of("address").map_values()
+
+        Returns:
+            A new `Expression` representing the values of the map.
+        """
+        return FunctionExpression("map_values", [self])
+
+    @expose_as_static
+    def map_entries(self) -> "Expression":
+        """Creates an expression that returns the entries of a map as an array of maps,
+        where each map contains a `"k"` property for the key and a `"v"` property for the value.
+        For example: `[{ "k": "key1", "v": "value1" }, ...]`.
+
+        Note:
+            While the backend generally preserves insertion order, relying on the
+            order of the output array is not guaranteed and should be avoided.
+
+        Example:
+            >>> Map({"city": "London", "country": "UK"}).map_entries()
+            >>> Field.of("address").map_entries()
+
+        Returns:
+            A new `Expression` representing the entries of the map.
+        """
+        return FunctionExpression("map_entries", [self])
+
     def regex_find(self, pattern: str | Constant[str] | Expression) -> "Expression":
         """Creates an expression that returns the first substring of a string expression that
         matches a specified regular expression.

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

@@ -107,6 +107,67 @@ tests:
                 title:
                   fieldReferenceValue: title
           name: select
+  - description: testMapSet
+    pipeline:
+      - Collection: books
+      - Where:
+          - FunctionExpression.equal:
+              - Field: title
+              - Constant: "Dune"
+      - Select:
+        - AliasedExpression:
+            - FunctionExpression.map_set:
+                - Field: awards
+                - "new_award"
+                - Constant: true
+            - "awards_set"
+    assert_results:
+      - awards_set:
+          hugo: true
+          nebula: true
+          new_award: true
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args:
+          - functionValue:
+              args:
+              - fieldReferenceValue: title
+              - stringValue: "Dune"
+              name: equal
+          name: where
+        - args:
+          - mapValue:
+              fields:
+                awards_set:
+                  functionValue:
+                    name: map_set
+                    args:
+                    - fieldReferenceValue: awards
+                    - stringValue: "new_award"
+                    - booleanValue: true
+          name: select
+  - description: testMapSetNone
+    pipeline:
+      - Collection: books
+      - Where:
+          - FunctionExpression.equal:
+              - Field: title
+              - Constant: "Dune"
+      - Select:
+        - AliasedExpression:
+            - FunctionExpression.map_set:
+                - Field: awards
+                - "hugo"
+                - Constant: null
+            - "awards_set"
+    assert_results:
+      - awards_set:
+          hugo: null
+          nebula: true
   - description: testMapRemove
     pipeline:
       - Collection: books
@@ -267,3 +328,119 @@ tests:
           a: "orig"
           b: "new"
           c: "new"
+  - description: testMapKeys
+    pipeline:
+      - Collection: books
+      - Where:
+          - FunctionExpression.equal:
+              - Field: title
+              - Constant: "Dune"
+      - Select:
+        - AliasedExpression:
+            - FunctionExpression.map_keys:
+                - Field: awards
+            - "award_keys"
+    assert_results:
+      - award_keys:
+          - hugo
+          - nebula
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args:
+          - functionValue:
+              args:
+              - fieldReferenceValue: title
+              - stringValue: "Dune"
+              name: equal
+          name: where
+        - args:
+          - mapValue:
+              fields:
+                award_keys:
+                  functionValue:
+                    name: map_keys
+                    args:
+                    - fieldReferenceValue: awards
+          name: select
+  - description: testMapValues
+    pipeline:
+      - Collection: books
+      - Where:
+          - FunctionExpression.equal:
+              - Field: title
+              - Constant: "Dune"
+      - Select:
+        - AliasedExpression:
+            - FunctionExpression.map_values:
+                - Field: awards
+            - "award_values"
+    assert_results:
+      - award_values:
+          - true
+          - true
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args:
+          - functionValue:
+              args:
+              - fieldReferenceValue: title
+              - stringValue: "Dune"
+              name: equal
+          name: where
+        - args:
+          - mapValue:
+              fields:
+                award_values:
+                  functionValue:
+                    name: map_values
+                    args:
+                    - fieldReferenceValue: awards
+          name: select
+  - description: testMapEntries
+    pipeline:
+      - Collection: books
+      - Where:
+          - FunctionExpression.equal:
+              - Field: title
+              - Constant: "Dune"
+      - Select:
+        - AliasedExpression:
+            - FunctionExpression.map_entries:
+                - Field: awards
+            - "award_entries"
+    assert_results:
+      - award_entries:
+          - k: hugo
+            v: true
+          - k: nebula
+            v: true
+    assert_proto:
+      pipeline:
+        stages:
+        - args:
+          - referenceValue: /books
+          name: collection
+        - args:
+          - functionValue:
+              args:
+              - fieldReferenceValue: title
+              - stringValue: "Dune"
+              name: equal
+          name: where
+        - args:
+          - mapValue:
+              fields:
+                award_entries:
+                  functionValue:
+                    name: map_entries
+                    args:
+                    - fieldReferenceValue: awards
+          name: select

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

@@ -1115,6 +1115,17 @@ def test_map_get(self):
         infix_instance = arg1.map_get(Constant.of(arg2))
         assert infix_instance == instance
 
+    def test_map_set(self):
+        arg1 = self._make_arg("Map")
+        arg2 = "key"
+        arg3 = "value"
+        instance = Expression.map_set(arg1, arg2, arg3)
+        assert instance.name == "map_set"
+        assert instance.params == [arg1, Constant.of(arg2), Constant.of(arg3)]
+        assert repr(instance) == "Map.map_set(Constant.of('key'), Constant.of('value'))"
+        infix_instance = arg1.map_set(Constant.of(arg2), arg3)
+        assert infix_instance == instance
+
     def test_map_remove(self):
         arg1 = self._make_arg("Map")
         arg2 = "key"
@@ -1136,6 +1147,33 @@ def test_map_merge(self):
         infix_instance = arg1.map_merge(arg2, arg3)
         assert infix_instance == instance
 
+    def test_map_keys(self):
+        arg1 = self._make_arg("Map")
+        instance = Expression.map_keys(arg1)
+        assert instance.name == "map_keys"
+        assert instance.params == [arg1]
+        assert repr(instance) == "Map.map_keys()"
+        infix_instance = arg1.map_keys()
+        assert infix_instance == instance
+
+    def test_map_values(self):
+        arg1 = self._make_arg("Map")
+        instance = Expression.map_values(arg1)
+        assert instance.name == "map_values"
+        assert instance.params == [arg1]
+        assert repr(instance) == "Map.map_values()"
+        infix_instance = arg1.map_values()
+        assert infix_instance == instance
+
+    def test_map_entries(self):
+        arg1 = self._make_arg("Map")
+        instance = Expression.map_entries(arg1)
+        assert instance.name == "map_entries"
+        assert instance.params == [arg1]
+        assert repr(instance) == "Map.map_entries()"
+        infix_instance = arg1.map_entries()
+        assert infix_instance == instance
+
     def test_mod(self):
         arg1 = self._make_arg("Left")
         arg2 = self._make_arg("Right")