Use orjson to improve JSON marshalling performance

https://github.com/ijl/orjson

Author: amotl

CHANGES.rst

@@ -5,6 +5,7 @@ Changes for crate
 Unreleased
 ==========
 
+- Use ``orjson`` to improve JSON marshalling performance. Thanks, @widmogrod.
 
 2024/11/23 1.0.1
 ================

setup.py

@@ -54,6 +54,7 @@ def read(path):
     packages=find_namespace_packages("src"),
     package_dir={"": "src"},
     install_requires=[
+        "orjson<4",
         "urllib3",
         "verlib2",
     ],

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(
@@ -180,7 +204,7 @@ def close(self):
 
 def _json_from_response(response):
     try:
-        return json.loads(response.data.decode("utf-8"))
+        return orjson.loads(response.data)
     except ValueError as ex:
         raise ProgrammingError(
             "Invalid server response of content-type '{}':\n{}".format(
@@ -223,7 +247,7 @@ def _raise_for_status_real(response):
     if response.status == 503:
         raise ConnectionError(message)
     if response.headers.get("content-type", "").startswith("application/json"):
-        data = json.loads(response.data.decode("utf-8"))
+        data = orjson.loads(response.data)
         error = data.get("error", {})
         error_trace = data.get("error_trace", None)
         if "results" in data:
@@ -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(