diff --git a/google/cloud/firestore_v1/pipeline_expressions.py b/google/cloud/firestore_v1/pipeline_expressions.py index 7a90dea1d..439b224cc 100644 --- a/google/cloud/firestore_v1/pipeline_expressions.py +++ b/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 diff --git a/tests/system/pipeline_e2e/array.yaml b/tests/system/pipeline_e2e/array.yaml index d32491d8b..acdded36b 100644 --- a/tests/system/pipeline_e2e/array.yaml +++ b/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 \ No newline at end of file diff --git a/tests/system/pipeline_e2e/logical.yaml b/tests/system/pipeline_e2e/logical.yaml index 0bc889ac1..bbb71921b 100644 --- a/tests/system/pipeline_e2e/logical.yaml +++ b/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 diff --git a/tests/unit/v1/test_pipeline_expressions.py b/tests/unit/v1/test_pipeline_expressions.py index 1546dbe66..84eb6cfe9 100644 --- a/tests/unit/v1/test_pipeline_expressions.py +++ b/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):