Skip to content

Commit

Permalink
fix: optimized protobuf access for performance (#155)
Browse files Browse the repository at this point in the history
More efficiently uses proto-plus wrappers, as well as inner protobuf attribute access, to greatly reduce the performance costs seen in version 2.0.0 (which stemmed from the introduction of proto-plus).

The size of the performance improvement scales with the number of attributes on each Entity, but in general, speeds once again closely approximate those from 1.15.

Fixes #145
Fixes #150
  • Loading branch information
craiglabenz authored Apr 20, 2021
1 parent 31bb053 commit 5b67daa
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 19 deletions.
43 changes: 24 additions & 19 deletions google/cloud/datastore/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,16 +49,21 @@ def _get_meaning(value_pb, is_list=False):
"""
meaning = None
if is_list:

values = (
value_pb._pb.array_value.values
if hasattr(value_pb, "_pb")
else value_pb.array_value.values
)

# An empty list will have no values, hence no shared meaning
# set among them.
if len(value_pb.array_value.values) == 0:
if len(values) == 0:
return None

# We check among all the meanings, some of which may be None,
# the rest which may be enum/int values.
all_meanings = [
_get_meaning(sub_value_pb) for sub_value_pb in value_pb.array_value.values
]
all_meanings = [_get_meaning(sub_value_pb) for sub_value_pb in values]
unique_meanings = set(all_meanings)
if len(unique_meanings) == 1:
# If there is a unique meaning, we preserve it.
Expand Down Expand Up @@ -119,11 +124,8 @@ def entity_from_protobuf(pb):
:rtype: :class:`google.cloud.datastore.entity.Entity`
:returns: The entity derived from the protobuf.
"""

if not getattr(pb, "_pb", False):
# Coerce raw pb type into proto-plus pythonic type.
proto_pb = entity_pb2.Entity(pb)
pb = pb
if not isinstance(pb, entity_pb2.Entity):
proto_pb = entity_pb2.Entity.wrap(pb)
else:
proto_pb = pb
pb = pb._pb
Expand Down Expand Up @@ -152,7 +154,7 @@ def entity_from_protobuf(pb):
if is_list and len(value) > 0:
exclude_values = set(
value_pb.exclude_from_indexes
for value_pb in value_pb.array_value.values
for value_pb in value_pb._pb.array_value.values
)
if len(exclude_values) != 1:
raise ValueError(
Expand Down Expand Up @@ -402,33 +404,36 @@ def _get_value_from_value_pb(value):
"""
if not getattr(value, "_pb", False):
# Coerce raw pb type into proto-plus pythonic type.
value = entity_pb2.Value(value)
value = entity_pb2.Value.wrap(value)

value_type = value._pb.WhichOneof("value_type")

if value_type == "timestamp_value":
# Do not access `._pb` here, as that returns a Timestamp proto,
# but this should return a Pythonic `DatetimeWithNanoseconds` value,
# which is found at `value.timestamp_value`
result = value.timestamp_value

elif value_type == "key_value":
result = key_from_protobuf(value.key_value)
result = key_from_protobuf(value._pb.key_value)

elif value_type == "boolean_value":
result = value.boolean_value
result = value._pb.boolean_value

elif value_type == "double_value":
result = value.double_value
result = value._pb.double_value

elif value_type == "integer_value":
result = value.integer_value
result = value._pb.integer_value

elif value_type == "string_value":
result = value.string_value
result = value._pb.string_value

elif value_type == "blob_value":
result = value.blob_value
result = value._pb.blob_value

elif value_type == "entity_value":
result = entity_from_protobuf(value.entity_value)
result = entity_from_protobuf(value._pb.entity_value)

elif value_type == "array_value":
result = [
Expand All @@ -437,7 +442,7 @@ def _get_value_from_value_pb(value):

elif value_type == "geo_point_value":
result = GeoPoint(
value.geo_point_value.latitude, value.geo_point_value.longitude,
value._pb.geo_point_value.latitude, value._pb.geo_point_value.longitude,
)

elif value_type == "null_value":
Expand Down
9 changes: 9 additions & 0 deletions tests/unit/test_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,15 @@ def test_entity_no_key(self):
self.assertIsNone(entity.key)
self.assertEqual(dict(entity), {})

def test_pb2_entity_no_key(self):
from google.cloud.datastore_v1.types import entity as entity_pb2

entity_pb = entity_pb2.Entity()
entity = self._call_fut(entity_pb)

self.assertIsNone(entity.key)
self.assertEqual(dict(entity), {})

def test_entity_with_meaning(self):
from google.cloud.datastore_v1.types import entity as entity_pb2
from google.cloud.datastore.helpers import _new_value_pb
Expand Down

0 comments on commit 5b67daa

Please sign in to comment.