Skip to content

Commit e45a8f7

Browse files
committed
Optimize Python C extension readStruct for nested structs
In this area of code, there already exists an optimization where we skip calling `klass(**kwargs)` on the top-level struct and instead init an empty object then PyObject_SetAttr the properties onto it.

This was done because Python's keyword argument matching is O(n) string comparisons per argument, making this expensive for structs with many fields. However, nested structs which are decoded within the C extension did not get this same optimization, causing them to go through the slow class initialization path. In this change, for mutable nested structs, we create the instance up front with a no-arg constructor (klass()) and set decoded fields directly via PyObject_SetAttr (the same way we fast-path the top-level struct). Immutable structs (TFrozenBase subclasses) continue to use the kwargs path since their generated __setattr__ blocks attribute mutation. Immutability is detected via PyObject_IsSubclass against TFrozenBase, with the class reference cached in a function-local static. Benchmarks show up to 3.6x speedup for deeply nested struct hierarchies, with no impact on flat structs. I had Claude generate me some benchmarks to showcase performance in different nested scenarios.

Note the test change. I asked about potentially changing this to look more like what the codegen produces here https://github.com/apache/thrift/pull/3349/changes#r3026288637. Currently, this fails because its not marked as frozen but has a erroring setattr
1 parent 197da84 commit e45a8f7

1 file changed

Lines changed: 35 additions & 6 deletions

File tree

lib/py/src/ext/protocol.tcc

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -874,17 +874,46 @@ PyObject* ProtocolBase<Impl>::decodeValue(TType type, PyObject* typeargs) {
874874
template <typename Impl>
875875
PyObject* ProtocolBase<Impl>::readStruct(PyObject* output, PyObject* klass, PyObject* spec_seq) {
876876
int spec_seq_len = PyTuple_Size(spec_seq);
877-
bool immutable = output == Py_None;
877+
bool immutable = false;
878878
ScopedPyObject kwargs;
879+
ScopedPyObject created_output;
879880
if (spec_seq_len == -1) {
880881
return nullptr;
881882
}
882883

883-
if (immutable) {
884-
kwargs.reset(PyDict_New());
885-
if (!kwargs) {
886-
PyErr_SetString(PyExc_TypeError, "failed to prepare kwargument storage");
887-
return nullptr;
884+
if (output == Py_None) {
885+
static PyObject* TBaseModule = nullptr;
886+
static PyObject* TFrozenBase = nullptr;
887+
if (!TFrozenBase) {
888+
if (!TBaseModule) {
889+
TBaseModule = PyImport_ImportModule("thrift.protocol.TBase");
890+
}
891+
if (!TBaseModule) {
892+
return nullptr;
893+
}
894+
TFrozenBase = PyObject_GetAttrString(TBaseModule, "TFrozenBase");
895+
if (!TFrozenBase) {
896+
return nullptr;
897+
}
898+
}
899+
// Immutable structs are produced by two codegen paths:
900+
// 1. "frozen2" mode: classes inherit from TFrozenBase
901+
// 2. "python.immutable" annotation: classes get a __setattr__ that raises TypeError
902+
immutable = PyObject_IsSubclass(klass, TFrozenBase)
903+
|| reinterpret_cast<PyTypeObject*>(klass)->tp_setattro != PyObject_GenericSetAttr;
904+
905+
if (immutable) {
906+
kwargs.reset(PyDict_New());
907+
if (!kwargs) {
908+
PyErr_SetString(PyExc_TypeError, "failed to prepare kwargument storage");
909+
return nullptr;
910+
}
911+
} else {
912+
created_output.reset(PyObject_CallObject(klass, nullptr));
913+
if (!created_output) {
914+
return nullptr;
915+
}
916+
output = created_output.get();
888917
}
889918
}
890919

0 commit comments

Comments
 (0)