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:
Tony Liao 2026-08-25 07:03:34 -07:00 committed by Copybara-Service
parent 7f86060aae
commit 9d98477ca9
2 changed files with 176 additions and 18 deletions

View file

@ -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()

View file

@ -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: