Start using orjson to improve JSON marshalling performance

amotl · amotl · commit c217a239cbac · 2025-01-15T23:58:38.000+01:00
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -5,6 +5,8 @@ Changes for crate
 Unreleased
 ==========
 
+- Started using ``orjson`` to improve JSON marshalling performance.
+  Thanks, @widmogrod.
 
 2024/11/23 1.0.1
 ================
diff --git a/setup.py b/setup.py
@@ -54,6 +54,7 @@ def read(path):
     packages=find_namespace_packages("src"),
     package_dir={"": "src"},
     install_requires=[
+        "orjson<4",
         "urllib3",
         "verlib2",
     ],
diff --git a/src/crate/client/http.py b/src/crate/client/http.py
@@ -37,6 +37,7 @@
 from urllib.parse import urlparse
 from uuid import UUID
 
+import orjson
 import urllib3
 from urllib3 import connection_from_url
 from urllib3.connection import HTTPConnection
@@ -107,6 +108,29 @@ def default(self, o):
         return json.JSONEncoder.default(self, o)
 
 
+def cratedb_json_encoder(obj):
+    """
+    Encoder function for orjson.
+
+    https://github.com/ijl/orjson#default
+    https://github.com/ijl/orjson#opt_passthrough_datetime
+    """
+    if isinstance(obj, (Decimal, UUID)):
+        return str(obj)
+    if isinstance(obj, datetime):
+        if obj.tzinfo is not None:
+            delta = obj - CrateJsonEncoder.epoch_aware
+        else:
+            delta = obj - CrateJsonEncoder.epoch_naive
+        return int(
+            delta.microseconds / 1000.0
+            + (delta.seconds + delta.days * 24 * 3600) * 1000.0
+        )
+    if isinstance(obj, date):
+        return calendar.timegm(obj.timetuple()) * 1000
+    return obj
+
+
 class Server:
     def __init__(self, server, **pool_kw):
         socket_options = _get_socket_opts(
@@ -334,7 +358,11 @@ def _create_sql_payload(stmt, args, bulk_args):
         data["args"] = args
     if bulk_args:
         data["bulk_args"] = bulk_args
-    return json.dumps(data, cls=CrateJsonEncoder)
+    return orjson.dumps(
+        data,
+        default=cratedb_json_encoder,
+        option=orjson.OPT_PASSTHROUGH_DATETIME,
+    )
 
 
 def _get_socket_opts(
