diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py index 49f3c0a97463..9f38fbf7db52 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_expressions.py +++ b/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. diff --git a/packages/google-cloud-firestore/tests/system/pipeline_e2e/map.yaml b/packages/google-cloud-firestore/tests/system/pipeline_e2e/map.yaml index 3e5e5de12e9d..0338c570b63e 100644 --- a/packages/google-cloud-firestore/tests/system/pipeline_e2e/map.yaml +++ b/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 diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py index 51c707528bc6..d673a2058bed 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_expressions.py +++ b/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")