Skip to content

Commit

Permalink
Performance optimization for upb python extend repeated scalars.
Browse files Browse the repository at this point in the history
Before:
_bench_append, 100000, 241.9
_bench_extend, 100000, 112.3
_bench_assign, 100000, 83.3
_bench_pybind11, 100000, 16.2

After:
_bench_append, 100000, 224.3
_bench_extend, 100000, 71.2
_bench_assign, 100000, 81.1
_bench_pybind11, 100000, 14.7
PiperOrigin-RevId: 711909495
  • Loading branch information
anandolee authored and copybara-github committed Jan 4, 2025
1 parent 7b71e08 commit 6c3ff5a
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 7 deletions.
9 changes: 9 additions & 0 deletions python/google/protobuf/internal/message_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -1058,6 +1058,13 @@ def testExtendStringWithNothing(self, message_module):
m.repeated_string.extend(empty_value)
self.assertSequenceEqual([], m.repeated_string)

def testExtendWithNoLen(self, message_module):
""" Test extending repeated fields with iterables but no len"""
m = message_module.TestAllTypes()
self.assertSequenceEqual([], m.repeated_int32)
m.repeated_int32.extend(i for i in range(2))
self.assertSequenceEqual([0, 1], m.repeated_int32)

def testExtendInt32WithPythonList(self, message_module):
"""Test extending repeated int32 fields with python lists."""
m = message_module.TestAllTypes()
Expand All @@ -1068,6 +1075,8 @@ def testExtendInt32WithPythonList(self, message_module):
self.assertSequenceEqual([0, 1, 2], m.repeated_int32)
m.repeated_int32.extend([3, 4])
self.assertSequenceEqual([0, 1, 2, 3, 4], m.repeated_int32)
with self.assertRaises(TypeError):
m.repeated_int32.extend([5, 6, 'hi', 7])

def testExtendFloatWithPythonList(self, message_module):
"""Test extending repeated float fields with python lists."""
Expand Down
31 changes: 24 additions & 7 deletions python/repeated.c
Original file line number Diff line number Diff line change
Expand Up @@ -181,15 +181,32 @@ PyObject* PyUpb_RepeatedContainer_Extend(PyObject* _self, PyObject* value) {
bool submsg = upb_FieldDef_IsSubMessage(f);
PyObject* e;

while ((e = PyIter_Next(it))) {
PyObject* ret;
if (submsg) {
ret = PyUpb_RepeatedCompositeContainer_Append(_self, e);
if (submsg) {
while ((e = PyIter_Next(it))) {
PyObject* ret = PyUpb_RepeatedCompositeContainer_Append(_self, e);
Py_XDECREF(ret);
Py_DECREF(e);
}
} else {
upb_Arena* arena = PyUpb_Arena_Get(self->arena);
Py_ssize_t size = PyObject_Size(value);
if (size < 0) {
// Some iterables may not have len. Size() will return -1 and
// set an error in such cases.
PyErr_Clear();
} else {
ret = PyUpb_RepeatedScalarContainer_Append(_self, e);
upb_Array_Reserve(arr, start_size + size, arena);
}
while ((e = PyIter_Next(it))) {
upb_MessageValue msgval;
if (!PyUpb_PyToUpb(e, f, &msgval, arena)) {
assert(PyErr_Occurred());
Py_DECREF(e);
break;
}
upb_Array_Append(arr, msgval, arena);
Py_DECREF(e);
}
Py_XDECREF(ret);
Py_DECREF(e);
}

Py_DECREF(it);
Expand Down

0 comments on commit 6c3ff5a

Please sign in to comment.