8000 Use `orjson` to improve JSON marshalling performance · crate/crate-python@c92e27b · GitHub
[go: up one dir, main page]
Skip to content

Commit c92e27b

Browse files
amotlamotl
committed
Use orjson to improve JSON marshalling performance
https://github.com/ijl/orjson
1 parent a2aae9b commit c92e27b

File tree

4 files changed

+40
-28
lines changed

4 files changed

+40
-28
lines changed

‎CHANGES.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ Changes for crate
55
Unreleased
66
==========
77

8+
- Use ``orjson`` to improve JSON marshalling performance. Thanks, @widmogrod.
89

910
2024/11/23 1.0.1
1011
================

‎setup.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ def read(path):
5454
packages=find_namespace_packages("src"),
5555
package_dir={"": "src"},
5656
install_requires=[
57+
"orjson<4",
5758
"urllib3",
5859
"verlib2",
5960
],

‎src/crate/client/http.py

Lines changed: 33 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import calendar
2424
import heapq
2525
import io
26-
import json
2726
import logging
2827
import os
2928
import re
@@ -37,6 +36,7 @@
3736
from urllib.parse import urlparse
3837
from uuid import UUID
3938

39+
import orjson
4040
import urllib3
4141
from urllib3 import connection_from_url
4242
from urllib3.connection import HTTPConnection
@@ -86,25 +86,31 @@ def super_len(o):
8686
return None
8787

8888

89-
class CrateJsonEncoder(json.JSONEncoder):
90-
epoch_aware = datetime(1970, 1, 1, tzinfo=timezone.utc)
91-
epoch_naive = datetime(1970, 1, 1)
92-
93-
def default(self, o):
94-
if isinstance(o, (Decimal, UUID)):
95-
return str(o)
96-
if isinstance(o, datetime):
97-
if o.tzinfo is not None:
98-
delta = o - self.epoch_aware
99-
else:
100-
delta = o - self.epoch_naive
101-
return int(
102-
delta.microseconds / 1000.0
103-
+ (delta.seconds + delta.days * 24 * 3600) * 1000.0
104-
)
105-
if isinstance(o, date):
106-
return calendar.timegm(o.timetuple()) * 1000
107-
return json.JSONEncoder.default(self, o)
89+
epoch_aware = datetime(1970, 1, 1, tzinfo=timezone.utc)
90+
epoch_naive = datetime(1970, 1, 1)
91+
92+
93+
def cratedb_json_encoder(obj):
94+
"""
95+
Encoder function for orjson.
96+
97+
https://github.com/ijl/orjson#default
98+
https://github.com/ijl/orjson#opt_passthrough_datetime
99+
"""
100+
if isinstance(obj, (Decimal, UUID)):
101+
return str(obj)
102+
if isinstance(obj, datetime):
103+
if obj.tzinfo is not None:
104+
delta = obj - epoch_aware
105+
else:
106+
delta = obj - epoch_naive
107+
return int(
108+
delta.microseconds / 1000.0
109+
+ (delta.seconds + delta.days * 24 * 3600) * 1000.0
110+
)
111+
if isinstance(obj, date):
112+
return calendar.timegm(obj.timetuple()) * 1000
113+
return obj
108114

109115

110116
class Server:
@@ -180,7 +186,7 @@ def close(self):
180186

181187
def _json_from_response(response):
182188
try:
183-
return json.loads(response.data.decode("utf-8"))
189+
return orjson.loads(response.data)
184190
except ValueError as ex:
185191
raise ProgrammingError(
186192
"Invalid server response of content-type '{}':\n{}".format(
@@ -223,7 +229,7 @@ def _raise_for_status_real(response):
223229
if response.status == 503:
224230
raise ConnectionError(message)
225231
if response.headers.get("content-type", "").startswith("application/json"):
226-
data = json.loads(response.data.decode("utf-8"))
232+
data = orjson.loads(response.data)
227233
error = data.get("error", {})
228234
error_trace = data.get("error_trace", None)
229235
if "results" in data:
@@ -334,7 +340,11 @@ def _create_sql_payload(stmt, args, bulk_args):
334340
data["args"] = args
335341
if bulk_args:
336342
data["bulk_args"] = bulk_args
337-
return json.dumps(data, cls=CrateJsonEncoder)
343+
return orjson.dumps(
344+
data,
345+
default=cratedb_json_encoder,
346+
option=orjson.OPT_PASSTHROUGH_DATETIME,
347+
)
338348

339349

340350
def _get_socket_opts(

‎tests/client/test_http.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -49,9 +49,9 @@
4949
)
5050
from crate.client.http import (
5151
Client,
52-
CrateJsonEncoder,
5352
_get_socket_opts,
5453
_remove_certs_for_non_https,
54+
cratedb_json_encoder,
5555
)
5656

5757
REQUEST = "crate.client.http.Server.request"
@@ -724,10 +724,10 @@ def test_username(self):
724724
class TestCrateJsonEncoder(TestCase):
725725
def test_naive_datetime(self):
726726
data = dt.datetime.fromisoformat("2023-06-26T09:24:00.123")
727-
result = json.dumps(data, cls=CrateJsonEncoder)
728-
self.assertEqual(result, "1687771440123")
727+
result = cratedb_json_encoder(data)
728+
self.assertEqual(result, 1687771440123)
729729

730730
def test_aware_datetime(self):
731731
data = dt.datetime.fromisoformat("2023-06-26T09:24:00.123+02:00")
732-
result = json.dumps(data, cls=CrateJsonEncoder)
733-
self.assertEqual(result, "1687764240123")
732+
result = cratedb_json_encoder(data)
733+
self.assertEqual(result, 1687764240123)

0 commit comments

Comments
 (0)
0