Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 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,69 @@ def find_nearest(
stages.FindNearest(field, vector, distance_measure, options)
)

def literals(self, *documents: Expression | dict) -> "_BasePipeline":
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.

Are you sure it makes sense to accept Expression here? Maybe it should just be dict | Map?

Expressions can evaluate to anything, and I don't know if it makes sense to have an int or bool Document

"""
Returns documents from a fixed set of predefined document objects.

This stage is commonly used for testing other stages in isolation,
though it can also be used as inputs to join conditions.
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.

This seems to conflict with the statement later:

The literals(...) stage can only be used as the first stage in a pipeline (or sub-pipeline)

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.


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)
Comment thread
Linchin marked this conversation as resolved.
Outdated
... .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
Comment thread
Linchin marked this conversation as resolved.
Outdated
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)
Comment thread
Linchin marked this conversation as resolved.
Outdated

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`.
Returns:
A new Pipeline object with this stage appended to the stage list.
"""
return self._append(stages.Literals(*documents))

def replace_with(
self,
field: Selectable,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,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: Expression | dict):
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,35 @@ 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.
- Constant:
value:
genre: "Science Fiction"
year: 1979
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.

This looks like it resolves to Constant({"genre": "Science Fiction", "year": 1979}). But our type hints forbid passing dicts as constants like this (mostly because Python's typing system isn't powerful enough to type nested datastructures, but still)

We provide a Map, which is supposed to fill that gap. Maybe you should test that out, and see if it would work as a declared type here?

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 @@ -417,6 +417,9 @@ def test_pipeline_execute_stream_equivalence():
("aggregate", (Field.of("n").as_("alias"),), stages.Aggregate),
("distinct", ("field_name",), stages.Distinct),
("distinct", (Field.of("n"), "second"), stages.Distinct),
("literals", (Field.of("a"),), stages.Literals),
("literals", ({"name": "joe"}, {"name": "bob"}), stages.Literals),
("literals", (Field.of("a"), {"name": "joe"}), stages.Literals),
],
)
def test_pipeline_methods(method, args, result_cls):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -517,6 +517,35 @@ 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


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