Skip to content

Commit 7538aa5

Browse files
LinchinLinchin
authored
feat: Add Rand and Trunc expressions (#16037)
1 parent 2158300 commit 7538aa5

File tree

3 files changed

+183
-0
lines changed

3 files changed

+183
-0
lines changed

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -405,6 +405,26 @@ def sqrt(self) -> "Expression":
405405
"""
406406
return FunctionExpression("sqrt", [self])
407407

408+
@expose_as_static
409+
def trunc(self, places: Expression | int | None = None) -> "Expression":
410+
"""Creates an expression that truncates the numeric value. If places is None,
411+
truncates to an integer. Otherwise, truncates the numeric value to the
412+
specified number of decimal places.
413+
414+
Example:
415+
>>> # Truncate the 'value' field to 2 decimal places.
416+
>>> Field.of("value").trunc(PipelineSource.literals(2))
417+
418+
Returns:
419+
A new `Expression` representing the truncated value.
420+
"""
421+
params = (
422+
[self, self._cast_to_expr_or_convert_to_constant(places)]
423+
if places is not None
424+
else [self]
425+
)
426+
return FunctionExpression("trunc", params)
427+
408428
@expose_as_static
409429
def logical_maximum(self, *others: Expression | CONSTANT_TYPE) -> "Expression":
410430
"""Creates an expression that returns the larger value between this expression
@@ -2109,3 +2129,15 @@ class CurrentTimestamp(FunctionExpression):
21092129

21102130
def __init__(self):
21112131
super().__init__("current_timestamp", [], use_infix_repr=False)
2132+
2133+
2134+
class Rand(FunctionExpression):
2135+
"""Creates an expression that generates a random number between 0.0 and 1.0 but not
2136+
including 1.0.
2137+
2138+
Returns:
2139+
A new `Expression` representing the rand operation.
2140+
"""
2141+
2142+
def __init__(self):
2143+
super().__init__("rand", [], use_infix_repr=False)

‎packages/google-cloud-firestore/tests/system/pipeline_e2e/math.yaml‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,89 @@ tests:
134134
- fieldReferenceValue: rating
135135
name: sqrt
136136
name: select
137+
- description: testTrunc
138+
pipeline:
139+
- Collection: books
140+
- Where:
141+
- FunctionExpression.equal_any:
142+
- Field: title
143+
- - Constant: "To Kill a Mockingbird" # rating 4.2
144+
- Constant: "Pride and Prejudice" # rating 4.5
145+
- Constant: "The Lord of the Rings" # rating 4.7
146+
- Select:
147+
- title
148+
- AliasedExpression:
149+
- FunctionExpression.trunc:
150+
- Field: rating
151+
- "trunc_rating"
152+
- AliasedExpression:
153+
- FunctionExpression.trunc:
154+
- Field: rating
155+
- FunctionExpression.add:
156+
- Constant: 0
157+
- Constant: 1
158+
- "trunc_rating_with_places"
159+
- Sort:
160+
- Ordering:
161+
- Field: title
162+
- ASCENDING
163+
assert_results:
164+
- title: "Pride and Prejudice"
165+
trunc_rating: 4.0
166+
trunc_rating_with_places: 4.5
167+
- title: "The Lord of the Rings"
168+
trunc_rating: 4.0
169+
trunc_rating_with_places: 4.7
170+
- title: "To Kill a Mockingbird"
171+
trunc_rating: 4.0
172+
trunc_rating_with_places: 4.2
173+
assert_proto:
174+
pipeline:
175+
stages:
176+
- args:
177+
- referenceValue: /books
178+
name: collection
179+
- args:
180+
- functionValue:
181+
args:
182+
- fieldReferenceValue: title
183+
- functionValue:
184+
args:
185+
- stringValue: "To Kill a Mockingbird"
186+
- stringValue: "Pride and Prejudice"
187+
- stringValue: "The Lord of the Rings"
188+
name: array
189+
name: equal_any
190+
name: where
191+
- args:
192+
- mapValue:
193+
fields:
194+
title:
195+
fieldReferenceValue: title
196+
trunc_rating:
197+
functionValue:
198+
args:
199+
- fieldReferenceValue: rating
200+
name: trunc
201+
trunc_rating_with_places:
202+
functionValue:
203+
args:
204+
- fieldReferenceValue: rating
205+
- functionValue:
206+
args:
207+
- integerValue: '0'
208+
- integerValue: '1'
209+
name: add
210+
name: trunc
211+
name: select
212+
- args:
213+
- mapValue:
214+
fields:
215+
direction:
216+
stringValue: ascending
217+
expression:
218+
fieldReferenceValue: title
219+
name: sort
137220
- description: testRoundFunctionExpressionessions
138221
pipeline:
139222
- Collection: books
@@ -306,4 +389,49 @@ tests:
306389
- fieldReferenceValue: rating
307390
- integerValue: '2'
308391
name: mod
392+
name: select
393+
- description: testRand
394+
pipeline:
395+
- Collection: books
396+
- Limit: 1
397+
- Select:
398+
- AliasedExpression:
399+
- And:
400+
- FunctionExpression.greater_than_or_equal:
401+
- Rand: []
402+
- Constant: 0.0
403+
- FunctionExpression.less_than:
404+
- Rand: []
405+
- Constant: 1.0
406+
- "is_valid_rand"
407+
assert_results:
408+
- is_valid_rand: true
409+
assert_proto:
410+
pipeline:
411+
stages:
412+
- args:
413+
- referenceValue: /books
414+
name: collection
415+
- args:
416+
- integerValue: '1'
417+
name: limit
418+
- args:
419+
- mapValue:
420+
fields:
421+
is_valid_rand:
422+
functionValue:
423+
name: and
424+
args:
425+
- functionValue:
426+
name: greater_than_or_equal
427+
args:
428+
- functionValue:
429+
name: rand
430+
- doubleValue: 0.0
431+
- functionValue:
432+
name: less_than
433+
args:
434+
- functionValue:
435+
name: rand
436+
- doubleValue: 1.0
309437
name: select

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1418,6 +1418,23 @@ def test_sqrt(self):
14181418
infix_instance = arg1.sqrt()
14191419
assert infix_instance == instance
14201420

1421+
def test_trunc(self):
1422+
arg1 = self._make_arg("Value")
1423+
instance = Expression.trunc(arg1)
1424+
assert instance.name == "trunc"
1425+
assert instance.params == [arg1]
1426+
assert repr(instance) == "Value.trunc()"
1427+
infix_instance = arg1.trunc()
1428+
assert infix_instance == instance
1429+
1430+
places = self._make_arg("Places")
1431+
instance_with_places = Expression.trunc(arg1, places)
1432+
assert instance_with_places.name == "trunc"
1433+
assert instance_with_places.params == [arg1, places]
1434+
assert repr(instance_with_places) == "Value.trunc(Places)"
1435+
infix_instance_with_places = arg1.trunc(places)
1436+
assert infix_instance_with_places == instance_with_places
1437+
14211438
def test_array_length(self):
14221439
arg1 = self._make_arg("Array")
14231440
instance = Expression.array_length(arg1)
@@ -1584,6 +1601,12 @@ def test_maximum(self):
15841601
infix_instance = arg1.maximum()
15851602
assert infix_instance == instance
15861603

1604+
def test_rand(self):
1605+
instance = expr.Rand()
1606+
assert instance.name == "rand"
1607+
assert instance.params == []
1608+
assert repr(instance) == "Rand()"
1609+
15871610
def test_array_agg(self):
15881611
arg1 = self._make_arg("Value")
15891612
instance = Expression.array_agg(arg1)

0 commit comments

Comments
 (0)