mirror of
https://github.com/protocolbuffers/protobuf
synced 2026-08-26 02:23:14 -04:00
Optimize pure Python parse path for custom JSON enum names.
This implementation is similar to our Java idea in cl/949670944, basically the idea is that we will have a local map (dictionary) that will be populated the very first time we try to parse into an enum field. The lifetime of this dictionary will be tied to the Parser instance, similar to the Java implementation. BENCHMARKS (baseline stats are recorded at cl/952170250): Python: ``` [BENCHMARK] ParseJsonDefault: med 10.06 us/op | p99 11.22 us/op | mean 10.10 ± 0.20 us/op [BENCHMARK] ParseJsonCustom: med 20.33 us/op | p99 22.23 us/op | mean 20.40 ± 0.34 us/op [BENCHMARK] ParseJsonUnknownIgnored: med 24.21 us/op | p99 26.25 us/op | mean 24.31 ± 0.44 us/op [BENCHMARK] ParseRepeatedJsonDefault: med 1.94 us/item | p99 2.05 us/item | mean 1.95 ± 0.03 us/item [BENCHMARK] ParseRepeatedJsonCustom: med 2.26 us/item | p99 2.45 us/item | mean 2.27 ± 0.03 us/item [BENCHMARK] ParseRepeatedJsonUnknownIgnored: med 5.31 us/item | p99 5.52 us/item | mean 5.32 ± 0.05 us/item ``` Cpp: ``` [BENCHMARK] ParseJsonDefault: med 8.66 us/op | p99 9.21 us/op | mean 8.70 ± 0.13 us/op [BENCHMARK] ParseJsonCustom: med 26.33 us/op | p99 29.71 us/op | mean 26.44 ± 0.69 us/op [BENCHMARK] ParseJsonUnknownIgnored: med 28.46 us/op | p99 31.80 us/op | mean 28.59 ± 0.67 us/op [BENCHMARK] ParseRepeatedJsonDefault: med 1.72 us/item | p99 2.19 us/item | mean 1.74 ± 0.06 us/item [BENCHMARK] ParseRepeatedJsonCustom: med 2.12 us/item | p99 2.38 us/item | mean 2.13 ± 0.04 us/item [BENCHMARK] ParseRepeatedJsonUnknownIgnored: med 4.61 us/item | p99 4.92 us/item | mean 4.62 ± 0.07 us/item ``` UPB: ``` [BENCHMARK] ParseJsonDefault: med 11.22 us/op | p99 12.29 us/op | mean 11.27 ± 0.21 us/op [BENCHMARK] ParseJsonCustom: med 37.54 us/op | p99 39.25 us/op | mean 37.65 ± 0.46 us/op [BENCHMARK] ParseJsonUnknownIgnored: med 41.18 us/op | p99 48.45 us/op | mean 41.30 ± 0.89 us/op [BENCHMARK] ParseRepeatedJsonDefault: med 2.10 us/item | p99 2.19 us/item | mean 2.10 ± 0.02 us/item [BENCHMARK] ParseRepeatedJsonCustom: med 2.58 us/item | p99 2.73 us/item | mean 2.59 ± 0.03 us/item [BENCHMARK] ParseRepeatedJsonUnknownIgnored: med 5.70 us/item | p99 5.90 us/item | mean 5.71 ± 0.04 us/item ``` Seeing these benchmarks, we can see a very noticeable improvement in `ParseRepeatedJsonCustom` and `ParseRepeatedJsonUnknownIgnored` cases -- which correspond to the pathological cases that we are worried about. Singular fields are taking longer to parse in our microbenchmark because the cache that we've added is tied to each Parser instance and must be re-instantiated across every `Parse` call. PiperOrigin-RevId: 970528683
This commit is contained in:
parent
7f86060aae
commit
9d98477ca9
2 changed files with 176 additions and 18 deletions
|
|
@ -2067,6 +2067,54 @@ class JsonFormatTest(JsonFormatBase):
|
|||
msg.armor, json_enumval_custom_string_pb2.Armor.ARMOR_UNKNOWN
|
||||
)
|
||||
|
||||
def testParseMapFieldWithCustomEnumNameOk(self):
|
||||
msg = json_enumval_custom_string_pb2.Knight()
|
||||
json_format.Parse(
|
||||
'{"armor_map": {"helmet": "gr8 helm", "boots": "sabaton", "chest": 8}}',
|
||||
msg,
|
||||
)
|
||||
self.assertEqual(
|
||||
msg.armor_map['helmet'],
|
||||
json_enumval_custom_string_pb2.Armor.ARMOR_GREAT_HELM,
|
||||
)
|
||||
self.assertEqual(
|
||||
msg.armor_map['boots'],
|
||||
json_enumval_custom_string_pb2.Armor.ARMOR_SABATON,
|
||||
)
|
||||
self.assertEqual(
|
||||
msg.armor_map['chest'],
|
||||
json_enumval_custom_string_pb2.Armor.ARMOR_HACHI_MAI_DO,
|
||||
)
|
||||
|
||||
def testGetCustomJsonEnumNames(self):
|
||||
enum_type = json_enumval_custom_string_pb2.Armor.DESCRIPTOR
|
||||
cache = {}
|
||||
names = json_format._GetCustomJsonEnumNames(enum_type, cache)
|
||||
self.assertIn('gr8 helm', names)
|
||||
self.assertEqual(names['gr8 helm'].name, 'ARMOR_GREAT_HELM')
|
||||
self.assertEqual(names['gr8 helm'].number, 1)
|
||||
self.assertIn('sabaton', names)
|
||||
self.assertEqual(names['sabaton'].number, 7)
|
||||
self.assertIn(names['sabaton'].name, ('ARMOR_SABATON', 'ARMOR_SOLLERET'))
|
||||
self.assertIn('8', names)
|
||||
self.assertEqual(names['8'].name, 'ARMOR_HACHI_MAI_DO')
|
||||
self.assertEqual(names['8'].number, 8)
|
||||
self.assertNotIn('ARMOR_GORGET', names)
|
||||
|
||||
# Verify cache is populated and subsequent calls return the cached dict.
|
||||
self.assertIn(enum_type, cache)
|
||||
self.assertIs(cache[enum_type], names)
|
||||
cached_names = json_format._GetCustomJsonEnumNames(enum_type, cache)
|
||||
self.assertIs(cached_names, names)
|
||||
|
||||
def testGetCustomJsonEnumNamesNoCustomOptions(self):
|
||||
enum_type = unittest_pb2.ForeignEnum.DESCRIPTOR
|
||||
cache = {}
|
||||
names = json_format._GetCustomJsonEnumNames(enum_type, cache)
|
||||
self.assertEqual(names, {})
|
||||
self.assertIn(enum_type, cache)
|
||||
self.assertEqual(cache[enum_type], {})
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
|
|||
|
|
@ -194,6 +194,32 @@ class _Printer(object):
|
|||
self.use_integers_for_enums = use_integers_for_enums
|
||||
self.descriptor_pool = descriptor_pool
|
||||
self.unquote_int64_if_possible = unquote_int64_if_possible
|
||||
self._enumvalue_json_extension = None
|
||||
|
||||
def _GetEnumValueJsonExtension(self):
|
||||
if self._enumvalue_json_extension is None:
|
||||
# Options are always put on the default pool, so we only search the
|
||||
# default pool.
|
||||
try:
|
||||
# Using reflection to FindExtensionByName is quite expensive, hence the
|
||||
# gymnastics to cache it into the instance attribute
|
||||
# _enumvalue_json_extension.
|
||||
# TODO: b/551998570 - Over the longer term, we can consider putting this
|
||||
# information in bootstrap files so that we don't have to rely on using
|
||||
# reflection to perform this lookup at all.
|
||||
self._enumvalue_json_extension = (
|
||||
descriptor_pool.Default().FindExtensionByName('pb.enumvalue.json')
|
||||
)
|
||||
except KeyError:
|
||||
self._enumvalue_json_extension = {}
|
||||
return self._enumvalue_json_extension or None
|
||||
|
||||
def _GetJsonEnumValueOption(self, ev):
|
||||
"""Helper to get the JsonEnumValueOptions for an enum value."""
|
||||
extension_descriptor = self._GetEnumValueJsonExtension()
|
||||
if extension_descriptor is None:
|
||||
return None
|
||||
return _GetJsonEnumValueOption(ev, extension_descriptor)
|
||||
|
||||
def ToJsonString(self, message, indent, sort_keys, ensure_ascii):
|
||||
js = self._MessageToJsonObject(message)
|
||||
|
|
@ -289,7 +315,7 @@ class _Printer(object):
|
|||
return None
|
||||
enum_value = field.enum_type.values_by_number.get(value, None)
|
||||
if enum_value is not None:
|
||||
option = _GetJsonEnumValueOption(enum_value)
|
||||
option = self._GetJsonEnumValueOption(enum_value)
|
||||
if option is not None:
|
||||
return option.string
|
||||
return enum_value.name
|
||||
|
|
@ -513,6 +539,20 @@ class _Parser(object):
|
|||
self.descriptor_pool = descriptor_pool
|
||||
self.max_recursion_depth = max_recursion_depth
|
||||
self.recursion_depth = 0
|
||||
self._custom_enum_names_cache = {}
|
||||
self._enumvalue_json_extension = None
|
||||
|
||||
def _GetEnumValueJsonExtension(self):
|
||||
if self._enumvalue_json_extension is None:
|
||||
# Options are always put on the default pool, so we only search the
|
||||
# default pool.
|
||||
try:
|
||||
self._enumvalue_json_extension = (
|
||||
descriptor_pool.Default().FindExtensionByName('pb.enumvalue.json')
|
||||
)
|
||||
except KeyError:
|
||||
self._enumvalue_json_extension = {}
|
||||
return self._enumvalue_json_extension or None
|
||||
|
||||
def ConvertMessage(self, value, message, path):
|
||||
"""Convert a JSON object into a message.
|
||||
|
|
@ -848,7 +888,12 @@ class _Parser(object):
|
|||
value_field = field.message_type.fields_by_name['value']
|
||||
for key in value:
|
||||
key_value = _ConvertScalarFieldValue(
|
||||
key, key_field, '{0}.key'.format(path), True
|
||||
key,
|
||||
key_field,
|
||||
'{0}.key'.format(path),
|
||||
self._custom_enum_names_cache,
|
||||
self._GetEnumValueJsonExtension(),
|
||||
require_str=True,
|
||||
)
|
||||
if value_field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_MESSAGE:
|
||||
self.ConvertMessage(
|
||||
|
|
@ -868,7 +913,13 @@ class _Parser(object):
|
|||
def _ConvertAndSetScalar(self, message, field, js_value, path):
|
||||
"""Convert scalar from js_value and assign it to message.field."""
|
||||
try:
|
||||
value = _ConvertScalarFieldValue(js_value, field, path)
|
||||
value = _ConvertScalarFieldValue(
|
||||
js_value,
|
||||
field,
|
||||
path,
|
||||
self._custom_enum_names_cache,
|
||||
self._GetEnumValueJsonExtension(),
|
||||
)
|
||||
if field.is_extension:
|
||||
message.Extensions[field] = value
|
||||
else:
|
||||
|
|
@ -884,7 +935,13 @@ class _Parser(object):
|
|||
repeated = message.Extensions[repeated_field]
|
||||
else:
|
||||
repeated = getattr(message, repeated_field.name)
|
||||
value = _ConvertScalarFieldValue(js_value, repeated_field, path)
|
||||
value = _ConvertScalarFieldValue(
|
||||
js_value,
|
||||
repeated_field,
|
||||
path,
|
||||
self._custom_enum_names_cache,
|
||||
self._GetEnumValueJsonExtension(),
|
||||
)
|
||||
repeated.append(value)
|
||||
except EnumStringValueParseError:
|
||||
if not self.ignore_unknown_fields:
|
||||
|
|
@ -900,39 +957,87 @@ class _Parser(object):
|
|||
js_value,
|
||||
map_field.message_type.fields_by_name['value'],
|
||||
path,
|
||||
self._custom_enum_names_cache,
|
||||
self._GetEnumValueJsonExtension(),
|
||||
)
|
||||
)
|
||||
except EnumStringValueParseError:
|
||||
if not self.ignore_unknown_fields:
|
||||
raise
|
||||
|
||||
def _GetJsonEnumValueOption(ev):
|
||||
|
||||
def _GetJsonEnumValueOption(ev, extension_descriptor):
|
||||
"""Helper to get the JsonEnumValueOptions for an enum value.
|
||||
|
||||
Args:
|
||||
ev: The EnumValueDescriptor.
|
||||
extension_descriptor: The extension descriptor for 'pb.enumvalue.json'.
|
||||
|
||||
Returns:
|
||||
The JsonEnumValueOptions message if the extension is present,
|
||||
otherwise None.
|
||||
"""
|
||||
try:
|
||||
extension_descriptor = descriptor_pool.Default().FindExtensionByName(
|
||||
'pb.enumvalue.json'
|
||||
)
|
||||
except KeyError:
|
||||
return None
|
||||
if ev.GetOptions().HasExtension(extension_descriptor):
|
||||
return ev.GetOptions().Extensions[extension_descriptor]
|
||||
return None
|
||||
|
||||
def _ConvertScalarFieldValue(value, field, path, require_str=False):
|
||||
|
||||
def _GetCustomJsonEnumNames(
|
||||
enum_type, custom_enum_names_cache, enumvalue_json_extension=None
|
||||
):
|
||||
"""Helper to get a mapping from custom JSON name to EnumValueDescriptor.
|
||||
|
||||
Args:
|
||||
enum_type: The EnumDescriptor.
|
||||
custom_enum_names_cache: A dict to store/lookup the cached map.
|
||||
enumvalue_json_extension: The extension descriptor for 'pb.enumvalue.json',
|
||||
or None to look it up in the default descriptor pool.
|
||||
|
||||
Returns:
|
||||
A dict mapping custom JSON name strings to EnumValueDescriptors.
|
||||
"""
|
||||
if enum_type in custom_enum_names_cache:
|
||||
return custom_enum_names_cache[enum_type]
|
||||
|
||||
custom_names = {}
|
||||
if enumvalue_json_extension is None:
|
||||
# Options are always put on the default pool, so we only search the default pool.
|
||||
try:
|
||||
enumvalue_json_extension = descriptor_pool.Default().FindExtensionByName(
|
||||
'pb.enumvalue.json'
|
||||
)
|
||||
except KeyError:
|
||||
enumvalue_json_extension = None
|
||||
|
||||
if enumvalue_json_extension is not None:
|
||||
for ev in enum_type.values:
|
||||
options = ev.GetOptions()
|
||||
if options.HasExtension(enumvalue_json_extension):
|
||||
option = options.Extensions[enumvalue_json_extension]
|
||||
if option.HasField('string'):
|
||||
custom_names[option.string] = ev
|
||||
|
||||
custom_enum_names_cache[enum_type] = custom_names
|
||||
return custom_names
|
||||
|
||||
|
||||
def _ConvertScalarFieldValue(
|
||||
value,
|
||||
field,
|
||||
path,
|
||||
custom_enum_names_cache,
|
||||
enumvalue_json_extension=None,
|
||||
require_str=False,
|
||||
):
|
||||
"""Convert a single scalar field value.
|
||||
|
||||
Args:
|
||||
value: A scalar value to convert the scalar field value.
|
||||
field: The descriptor of the field to convert.
|
||||
path: parent path to log parse error info.
|
||||
custom_enum_names_cache: A dict to store/lookup custom enum names.
|
||||
enumvalue_json_extension: The extension descriptor for 'pb.enumvalue.json',
|
||||
or None if not loaded.
|
||||
require_str: If True, the field value must be a str.
|
||||
|
||||
Returns:
|
||||
|
|
@ -968,12 +1073,17 @@ def _ConvertScalarFieldValue(value, field, path, require_str=False):
|
|||
# Convert an enum value.
|
||||
enum_value = field.enum_type.values_by_name.get(value, None)
|
||||
# First check to see if we have a custom enum string.
|
||||
if enum_value is None:
|
||||
for ev in field.enum_type.values:
|
||||
option = _GetJsonEnumValueOption(ev)
|
||||
if option is not None and option.string == value:
|
||||
enum_value = ev
|
||||
break
|
||||
if (
|
||||
enum_value is None
|
||||
and isinstance(value, str)
|
||||
and enumvalue_json_extension is not None
|
||||
):
|
||||
custom_names = _GetCustomJsonEnumNames(
|
||||
field.enum_type,
|
||||
custom_enum_names_cache,
|
||||
enumvalue_json_extension,
|
||||
)
|
||||
enum_value = custom_names.get(value, None)
|
||||
# If not, try parsing it as an integer.
|
||||
if enum_value is None:
|
||||
try:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue