`orjson` for much faster JSON operations?

[DanielYang59]

I was trying to optimize the performance for a high-throughput use case (proposed by @janosh) where one might need to load/save large quantities of Structure as JSON, and I'm currently starting with the saving phase of the workflow.

A quick line profiling reveals two significant bottlenecks:

• as_dict method of Structure: I did some preliminary profiling, and looks like this is owing to (unnecessarily) repeated calculation of Lattice properties like lengths/angles during serialization (and I believe proper caching could resolve this with correct cache voiding mechanism upon internal _matrix change). But this would be a separate PR so won't cover too much here.
• Topic of this thread: slow json.dump to file

Total time: 65.3412 s
File: show_structure_json_bottleneck.py
Function: save_to_file at line 38

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    38                                           @profile
    39                                           def save_to_file(tmpdir):
    40      1001        921.1      0.9      0.0      for i, struct in enumerate(STRUCTURES):
    41      1000   32742869.8  32742.9     50.1          d = struct.as_dict()
    42      1000       9134.1      9.1      0.0          path = os.path.join(tmpdir, f"struct_{i}.json")
    43      2000      77023.5     38.5      0.1          with open(path, "w") as f:
    44      1000   32511267.6  32511.3     49.8              json.dump(d, f)

---

And if we replace builtin json with the faster drop-in replacement orjson (I'm not aware of any caveat so far), we got (tested on WSL2 Ubuntu 22.04):

• almost 10x speed up converting Structure dict to string
• more than 10x speed up for saving Structure dict to file
• a little speed up loading JSON back

== Dump to String ==
    json: 0.9821 s
  orjson: 0.1173 s

== Save to File ==
  json_plain: 3.2918 s, 105549.10 KB
orjson_plain: 0.2284 s, 97192.09 KB
     json_gz: 8.0444 s, 32786.10 KB
   orjson_gz: 3.3784 s, 32217.06 KB

== Load from File ==
  json_plain: 3.1989 s
orjson_plain: 2.6992 s
     json_gz: 3.4490 s
   orjson_gz: 2.9699 s

Script: compare_structure_json.txt

As the speed up is quite obvious with almost no code change (also considering JSON is heavily used around the code base, so not only Structure would benefit from this), perhaps we could consider adding orjson as the default JSON handler? (I'm afraid adding it as an optional dependency would incur too much code change to detect whether orjson is available)

But also I'm a bit unsure about the migration process as current json.dumps takes kwargs (e.g. #4295) and this might cause breakage:

pymatgen/src/pymatgen/core/structure.py

Line 2963 in 1e36ee5

json_str = json.dumps(self.as_dict(), **kwargs)

monty JSON and lack of custom JSON encoder support

orjson is already an optional dependency for monty so maybe it's also good for us to double check if orjson is fully used from monty's side as quite some JSON-related workloads are implicitly handled by monty.

Also as orjson doesn't directly support custom encoder/decoder (e.g. MontyEncoder), some more thoughts might be needed to this end

[DanielYang59]

After posting this "proposal", I realized two issues:

• The saving/loading process is more or less detatched from the serialization (as_dict), so one could always install orjson locally without any change needed from pymatgen's side.
• For some processes where the json.dumps is integrated, very likely a custom JSON encoder/decoder is needed (e.g. MontyEncoder). Looks like orjson natively support serialization of np.array datetime and dataclass (but I haven't tested myself, not sure if results is consistent with MontyEncoder), also there is default to support arbitary type similar to the one from monty encoder:

  It serializes dataclass, datetime, numpy, and UUID instances natively.

But in any case I'm happy to give this a try

[shyuep]

While I don't have a problem with this in principle as long as orjson is done as an optional import, e.g.,

try:
    import orjson as json
except ImportError:
    import json

this only works when you don't need a custom encoder like MontyEncoder or MontyDecoder. This is severely limiting. Now, there are ways to get around that of course. But I think it is not worth the effort. People who really need to work that HT workflow can implement those hacks themselves (e.g., preconverting to dicts and then using orjson to do the dump). The majority of people are dumping a few structures and the difference between 2 ms and 10 ms is not going to be noticeable.

Note: I attempted to do this in Monty and I concluded it is too much complexity to be worth it.

[shyuep]

I would also advocate for a more practical view of optimizations. The most important optimizations are those that serve > 80% of users and result in at least a 10-20%, preferably order of magnitude speed improvements, with almost no backwards incompatibility. E.g., optimizing the nearest neighbor finding code or the recent PR by @kavanase that reduced vasprun parsing times AND simplified the code by removing Incar.get calls.

Anything that serves esoteric needs and/or result in small-ish speed increases is not worth our time, especially if code complexity has to significantly increase. Your code suggests that for 1000 structures, you are saving 1s of time. Very few people are going to work with anything more than 10000 structures and loading / dumping is a one time affair. For people who want to work with 1-100 million datasets, they have the skills to optimize for their own use case.

[kavanase]

The most important optimizations are those that serve > 80% of users and result in at least a 10-20%, preferably order of magnitude speed improvements, with almost no backwards incompatibility.

I definitely agree with this. If it's of any use to note; the biggest pymatgen runtime bottlenecks in our workflows (which I think is similar to most users) are Vasprun parsing (better now but even faster would be nice) and StructureMatcher utilities (which has been discussed here: #2593)

[DanielYang59]

The most important optimizations are those that serve > 80% of users and result in at least a 10-20%, preferably order of magnitude speed improvements, with almost no backwards incompatibility.

I totally agree with you that an "ideal" optimization should bring the most performance benefit to most user while adding minimal code complexity nor backwards incompatibility.

In this regard replacing json with orjson seem to tick most of the boxes in my opinion:

• Speed benefit (quote orjson doc):

  orjson.dumps() is something like 10x as fast as json,
  orjson.loads() is something like 2x as fast as json

• Applicability: I agree loading millions of Structure might be a niche need, but orjson would speed up many more JSON related workload across the code base. This is especially relevant given how extensively JSON is used in pymatgen, and considering pymatgen also serves as the backend for the Materials Project web application (if I understand correctly).
• Code change: for those don't need a custom encoder/decoder, this would largely be a drop-in replacement (assuming it could be a required dependency otherwise would require a lot of checking for orjson availability)

this only works when you don't need a custom encoder like MontyEncoder or MontyDecoder

Yes this is a big limiting factor for this migration and I haven't got much good thoughts on this (but still we could leave them as is):

• orjson natively support dataclass, datetime, numpy, and UUID, so if some instance only needs a custom encoder/decoder for these datatypes, we could just get rid of that

---

the biggest pymatgen runtime bottlenecks in our workflows (which I think is similar to most users) are Vasprun parsing (better now but even faster would be nice) and StructureMatcher utilities

Hi @kavanase thanks for pointing out the bottleneck in your workflow, this would be immensely helpful for people like me who want to tweak some performance but don't know where to get started :D

[shyuep]

I am aware json is being used in the code base. I am merely pointing out that the performance gains are small when your datasets are not large. For example, the No. 1 most read file is probably periodic_table.json.

%timeit -n 100 orjson.loads(open("periodic_table.json").read())
603 μs ± 46.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit -n 100 json.load(open("periodic_table.json"))
1.13 ms ± 22 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

On paper, 1.13 ms is definitely a 2x improvement over 603 us. Will anyone feel an improvement of 603 us? Probably not.

But this is one use case where you can easily do a drop-in support for orjson without too much issue.

[shyuep]

One last thing: a major issue I have with orjson is that they don't even support simple dump / load from IO streams, in addition lack of support custom encoders/decoders. That's actually very odd given how easy it is to support those. What it means is that we always have to dump to string and then write it separately.

If you want to find alternatives, here they are. https://catnotfoundnear.github.io/finding-the-fastest-python-json-library-on-all-python-versions-8-compared.html

Finally, feel free to add orjson support wherever it is going to be trivial. But again, I wouldn't spend too much time on this. I actually think there really aren't that many locations in the code where we are explicitly dumping json.

[shyuep]

@kavanase You will be pleased to know that with the recent commit, I sped up Vasprun parsing by ~33% using lxml.

Vasprun("tests/files/io/vasp/outputs/vasprun.r2scan.xml.gz")

which is quite small, took about 0.6s with Python's stdlib ElementTree, but 0.4s with lxml's ElementTree.

[kavanase]

Ah that would be very nice! I've wondered before if there was a faster version of XML ElementTree available somewhere...

I've just run a few tests, and unfortunately seem to find that for large vasprun files, this seems to decrease parsing speed.
Parsing the attached Vasprun (from a SOC defect supercell calculation) 10 times on my MacBook takes 38-39s with the previous code, but 44-48s with the new lxml code (running this parsing loop 5 times repeatedly to make sure).

vasprun.xml.gz

I tried some profiling with cProfile but lxml seems to block its recording to identify the bottleneck

[shyuep]

Hmm really? By all accounts, lxml should be faster. Let me investigate a bit.

[kavanase]

Yeah I was surprised, but those are the numbers I was consistently getting with that large vasprun file

[DanielYang59]

I just had a quick test on my WSL2 Ubuntu 22.04 machine in case it's some issue with specific platform:

• on 2025.4.24: Average over 5 runs (import + parse): 3.902 seconds
• on 2025.5.1: Average over 5 runs (import + parse): 6.887 seconds

I didn't do any profiling either so not sure where the regression comes from

Thanks for the helpful discussion in any case and enjoy the holiday :D

[kavanase]

When I tested, I checked out the specific lxml commit (before the ptable JSON stuff) so it shouldn't be that at least. lxml doesn't seem to play well with cProfile for some reason, as it just doesn't show what it's doing, whereas with XML ElementTree it showed the iterator, feed and read_events functions as taking the bulk of the IO time.

[shyuep]

@kavanase I found that setting remove_blank_text=True yielded a signfiicant increase in speed. Can you try the latest version on master.

[kavanase]

I tested again (with that vasprun shared above), and it's better. The lxml parsing time is now 35.5-38 seconds, versus 38-40s originally (and 44-48s with lxml but without remove_blank_text=True). So a bit faster than the original code now!