Skip to content
Merged
Show file tree
Hide file tree
Changes from 12 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,7 @@ def find_nearest(
stages.FindNearest(field, vector, distance_measure, options)
)


def replace_with(
self,
field: Selectable,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
from google.cloud.firestore_v1.base_document import BaseDocumentReference
from google.cloud.firestore_v1.base_query import BaseQuery
from google.cloud.firestore_v1.client import Client
from google.cloud.firestore_v1.pipeline_expressions import Expression


PipelineType = TypeVar("PipelineType", bound=_BasePipeline)
Expand Down Expand Up @@ -108,3 +109,66 @@ def documents(self, *docs: "BaseDocumentReference") -> PipelineType:
a new pipeline instance targeting the specified documents
"""
return self._create_pipeline(stages.Documents.of(*docs))

def literals(self, *documents: dict[str, "Expression" | CONSTANT_TYPE]) -> PipelineType:
"""
Returns documents from a fixed set of predefined document objects.

This stage is commonly used for testing other stages in isolation,
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@daniel-sanche I think what the documentation means is that, the Literals stage could serve as one of the inputs of a join operation, which is not contradictory to it being a first stage, so I think it is reasonable to keep this sentence and make it consistent with documentation.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd still prefer to take it out since we don't yet support the join stage. But if you think it seems to match the coming feature, we can probably leave it

though it can also be used as inputs to join conditions.

Example:
>>> from google.cloud.firestore_v1.pipeline_expressions import Constant
>>> documents = [
... {"name": "joe", "age": 10},
... {"name": "bob", "age": 30},
... {"name": "alice", "age": 40}
... ]
>>> pipeline = client.pipeline()
... .literals(*documents)
... .where(field("age").lessThan(35))

Output documents:
```json
[
{"name": "joe", "age": 10},
{"name": "bob", "age": 30}
]
```

Behavior:
The `literals(...)` stage can only be used as the first stage in a pipeline (or
sub-pipeline). The order of documents returned from the `literals` matches the
order in which they are defined.

While literal values are the most common, it is also possible to pass in
expressions, which will be evaluated and returned, making it possible to test
out different query / expression behavior without first needing to create some
test data.

For example, the following shows how to quickly test out the `length(...)`
function on some constant test sets:

Example:
>>> from google.cloud.firestore_v1.pipeline_expressions import Constant
>>> documents = [
... {"x": Constant.of("foo-bar-baz").char_length()},
... {"x": Constant.of("bar").char_length()}
... ]
>>> pipeline = client.pipeline().literals(*documents)

Output documents:
```json
[
{"x": 11},
{"x": 3}
]
```

Args:
*documents: One or more documents to be returned by this stage. Each can be a `dict`
or an `Expression`.
Comment thread
Linchin marked this conversation as resolved.
Outdated
Returns:
A new Pipeline object with this stage appended to the stage list.
"""
return self._create_pipeline(stages.Literals(*documents))
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
AggregateFunction,
AliasedExpression,
BooleanExpression,
CONSTANT_TYPE,
Expression,
Field,
Ordering,
Expand Down Expand Up @@ -342,6 +343,23 @@ def _pb_args(self):
return [Value(integer_value=self.limit)]


class Literals(Stage):
"""Returns documents from a fixed set of predefined document objects."""

def __init__(self, *documents: dict[str, "Expression" | CONSTANT_TYPE]):
super().__init__("literals")
self.documents = documents

def _pb_args(self):
args = []
for doc in self.documents:
if hasattr(doc, "_to_pb"):
args.append(doc._to_pb())
Comment thread
Linchin marked this conversation as resolved.
Outdated
else:
args.append(encode_value(doc))
return args


class Offset(Stage):
"""Skips a specified number of documents."""

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -684,4 +684,33 @@ tests:
- args:
- fieldReferenceValue: awards
- stringValue: full_replace
name: replace_with
name: replace_with
- description: literals
pipeline:
- Literals:
- title: "The Hitchhiker's Guide to the Galaxy"
author: "Douglas Adams"
Comment thread
daniel-sanche marked this conversation as resolved.
- genre: "Science Fiction"
year: 1979
assert_results:
- title: "The Hitchhiker's Guide to the Galaxy"
author: "Douglas Adams"
- genre: "Science Fiction"
year: 1979
assert_proto:
pipeline:
stages:
- args:
- mapValue:
fields:
author:
stringValue: "Douglas Adams"
title:
stringValue: "The Hitchhiker's Guide to the Galaxy"
- mapValue:
fields:
genre:
stringValue: "Science Fiction"
year:
integerValue: '1979'
name: literals
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should also have tests here that cover the different input types we support

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! I added additional type to test.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we still need more examples. Both of these are dict-like, so it seems like a pretty basic test. What if someone passes in Constant(1)? Or Constant("test").byte_length()? We say we support all expressions, how are non-dict types represented?

It would also be good to add some extra stages, to make sure this works like others

You can use gemini to create a few extra test scenarios. I usually don't include assert_proto on all of them, because it can be excessive

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Referring to the golang implementation, the only type accepted is a list of key-value pairs. So I don't think Constant(1) or Constant("test").byte_length() are in scope.

But I think it's a good idea to add more tests. I will also use mapValue instead of Constant.

Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,17 @@ def test_documents(self):
assert first_stage.paths[1] == "/a/2"
assert first_stage.paths[2] == "/a/3"

def test_literals(self):
from google.cloud.firestore_v1.pipeline_expressions import Field

instance = self._make_client().pipeline()
documents = (Field.of("a"), {"name": "joe"})
Comment thread
Linchin marked this conversation as resolved.
Outdated
ppl = instance.literals(*documents)
assert isinstance(ppl, self._expected_pipeline_type)
assert len(ppl.stages) == 1
first_stage = ppl.stages[0]
assert isinstance(first_stage, stages.Literals)


class TestPipelineSourceWithAsyncClient(TestPipelineSource):
"""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -517,6 +517,89 @@ def test_to_pb(self):
assert len(result.options) == 0


class TestLiterals:
def _make_one(self, *args, **kwargs):
return stages.Literals(*args, **kwargs)

def test_ctor(self):
val1 = Constant.of({"a": 1})
Comment thread
Linchin marked this conversation as resolved.
Outdated
val2 = {"b": 2}
instance = self._make_one(val1, val2)
assert instance.documents == (val1, val2)
assert instance.name == "literals"
Comment thread
daniel-sanche marked this conversation as resolved.

def test_repr(self):
val1 = Constant.of({"a": 1})
instance = self._make_one(val1, {"b": 2})
repr_str = repr(instance)
assert repr_str == "Literals(documents=(Constant.of({'a': 1}), {'b': 2}))"

def test_to_pb(self):
val1 = Constant.of({"a": 1})
val2 = {"b": 2}
instance = self._make_one(val1, val2)
result = instance._to_pb()
assert result.name == "literals"
assert len(result.args) == 2
assert result.args[0].map_value.fields["a"].integer_value == 1
assert result.args[1].map_value.fields["b"].integer_value == 2
assert len(result.options) == 0

def test_ctor_extended_types(self):
import datetime
from google.cloud.firestore_v1._helpers import GeoPoint
from google.cloud.firestore_v1.vector import Vector

doc = {
"a": 1,
"b": "string",
"c": 3.14,
"d": True,
"e": None,
"f": b"bytes",
"g": datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
"h": GeoPoint(1.0, 2.0),
"i": Vector([1.0, 2.0])
}
instance = self._make_one(doc)
assert instance.documents == (doc,)
assert instance.name == "literals"

def test_to_pb_extended_types(self):
import datetime
from google.cloud.firestore_v1._helpers import GeoPoint
from google.cloud.firestore_v1.vector import Vector

doc = {
"a": 1,
"b": "string",
"c": 3.14,
"d": True,
"e": None,
"f": b"bytes",
"g": datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
"h": GeoPoint(1.0, 2.0),
"i": Vector([1.0, 2.0])
}
instance = self._make_one(doc)
result = instance._to_pb()
assert result.name == "literals"
assert len(result.args) == 1

fields = result.args[0].map_value.fields
assert fields["a"].integer_value == 1
assert fields["b"].string_value == "string"
assert fields["c"].double_value == 3.14
assert fields["d"].boolean_value is True
assert fields["e"].null_value == 0
assert fields["f"].bytes_value == b"bytes"
assert fields["g"].timestamp_value == datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc)
assert fields["h"].geo_point_value.latitude == 1.0
assert fields["h"].geo_point_value.longitude == 2.0
assert fields["i"].map_value.fields["value"].array_value.values[0].double_value == 1.0
assert fields["i"].map_value.fields["value"].array_value.values[1].double_value == 2.0


class TestOffset:
def _make_one(self, *args, **kwargs):
return stages.Offset(*args, **kwargs)
Expand Down
Loading