From 76c471d713f8d1ba38e1d82359f4c6a58bff1189 Mon Sep 17 00:00:00 2001 From: Jonas Daniels Date: Wed, 28 Oct 2026 14:38:32 -0700 Subject: [PATCH] perf: send query responses up to 3x faster by reusing mapped ASTs and their JSON - An interned query keeps its server-named AST, so a query built again maps once. - An AST keeps its JSON text; the query endpoint joins each query's text instead of stringifying and scanning the whole response. - js_stringify_bytes finds array-index keys in one scan instead of ten. - The builder looks up a relationship's base query without a closure, and keys a column-to-scalar condition as one flat tuple. Co-Authored-By: Claude Opus 5.5 --- bench/cpu.py | 4 +- src/zero_server/integrations/starlette.py | 4 +- src/zero_server/protocol/ast.py | 19 +++-- src/zero_server/protocol/custom_queries.py | 6 ++ src/zero_server/protocol/query_server.py | 16 +++- src/zero_server/query/builder.py | 95 ++++++++++++++-------- src/zero_server/server/process_queries.py | 17 ++-- src/zero_server/shared/js.py | 6 +- tests/unit/test_protocol.py | 4 + tests/wire/test_query.py | 4 +- 10 files changed, 120 insertions(+), 55 deletions(-) diff --git a/bench/cpu.py b/bench/cpu.py index d2cecfb..fb20a0c 100644 --- a/bench/cpu.py +++ b/bench/cpu.py @@ -34,7 +34,7 @@ from zero_server.protocol.wire import parse_wire from zero_server.query.builder import create_builder from zero_server.server.process_mutations import handle_mutate_request -from zero_server.server.process_queries import handle_query_request +from zero_server.server.process_queries import handle_query_request, query_response_bytes from zero_server.server.server_schema import get_server_schema from zero_server.server.zql_database import ZQLDatabase from zero_server.z2s.compiler import compile_query @@ -57,7 +57,7 @@ async def query(body: Any = body) -> Any: async def query_json(text: bytes = text) -> bytes: response = await handle_query_request(handler=handler, schema=query_schema, body=text, user_id="u1") - return json_response(response.to_json()) + return query_response_bytes(response) out[f"query.{kind}.b{n}"] = query out[f"query_json.{kind}.b{n}"] = query_json diff --git a/src/zero_server/integrations/starlette.py b/src/zero_server/integrations/starlette.py index f1f0cf8..54a343b 100644 --- a/src/zero_server/integrations/starlette.py +++ b/src/zero_server/integrations/starlette.py @@ -20,7 +20,7 @@ from zero_server.query.registry import must_get_query from zero_server.schema import Schema from zero_server.server.process_mutations import Database, TransactFn, handle_mutate_request -from zero_server.server.process_queries import handle_query_request +from zero_server.server.process_queries import handle_query_request, query_response_bytes from zero_server.shared.js import js_stringify_bytes from zero_server.shared.json import JSONValue from zero_server.shared.registry import Registry @@ -88,7 +88,7 @@ async def query(request: Request) -> Response: body=await request.body(), user_id=user_id, ) - return _JSONResponse(response.to_json()) + return Response(query_response_bytes(response), media_type="application/json") routes.append(Route(query_path, query, methods=["POST"])) diff --git a/src/zero_server/protocol/ast.py b/src/zero_server/protocol/ast.py index a339b12..f8a181e 100644 --- a/src/zero_server/protocol/ast.py +++ b/src/zero_server/protocol/ast.py @@ -17,7 +17,7 @@ from zero_server.protocol.wire import dump_json, json_model, parse_wire from zero_server.schema import NameMapper as SchemaNameMapper from zero_server.shared.asserts import AssertionFailedError -from zero_server.shared.js import js_keys, js_sort +from zero_server.shared.js import js_keys, js_sort, js_stringify_bytes from zero_server.shared.json import JSONValue if TYPE_CHECKING: @@ -68,10 +68,10 @@ def _names(names: Sequence[Any]) -> list[Any]: # The nodes are plain slotted classes: the builder makes dozens per query, and a Pydantic model costs # several times as much to build. They are not frozen, which would cost as much again, but no code -# changes a node once built: every transformation builds new ones. `to_json` writes the keys in the order -# upstream's normalized AST has them, which every AST the builder makes is, and leaves out a field that is -# `None`. `AST.from_json` parses through the wire -# schema in `ast_wire`, which accepts exactly what upstream's `astSchema` accepts. +# changes a node once built: every transformation builds new ones, so an AST keeps its JSON text once +# written. `to_json` writes the keys in the order upstream's normalized AST has them, which every AST the +# builder makes is, and leaves out a field that is `None`. `AST.from_json` parses through the wire schema +# in `ast_wire`, which accepts exactly what upstream's `astSchema` accepts. @json_model @@ -231,6 +231,7 @@ class AST: start: Bound | None = None limit: int | float | None = None order_by: Ordering | None = None + _text: bytes | None = field(default=None, init=False, repr=False, compare=False) def _copy(self) -> AST: """A shallow copy, for a caller to change one field of. Twice as fast as `AST(...)` with every field.""" @@ -243,6 +244,7 @@ def _copy(self) -> AST: copy.start = self.start copy.limit = self.limit copy.order_by = self.order_by + copy._text = None return copy def to_json(self) -> dict[str, Any]: @@ -263,6 +265,13 @@ def to_json(self) -> dict[str, Any]: out["orderBy"] = [_names(part) for part in self.order_by] return out + def to_json_bytes(self) -> bytes: + """`to_json()` as the text `JSON.stringify` writes, kept for the next response that sends this AST.""" + text = self._text + if text is None: + text = self._text = js_stringify_bytes(self.to_json()) + return text + @classmethod def from_json(cls, value: object) -> AST: """Parse `value` off the wire, raw bytes or decoded JSON, as upstream's `astSchema` does. diff --git a/src/zero_server/protocol/custom_queries.py b/src/zero_server/protocol/custom_queries.py index 8a081cf..6c52184 100644 --- a/src/zero_server/protocol/custom_queries.py +++ b/src/zero_server/protocol/custom_queries.py @@ -11,6 +11,7 @@ from zero_server.protocol.ast import AST from zero_server.protocol.wire import PlainJSON, RefusesNullOffWire, WireModel, json_model +from zero_server.shared.js import js_stringify_bytes from zero_server.shared.json import JSONValue @@ -38,6 +39,11 @@ class TransformedQuery: def to_json(self) -> dict[str, Any]: return {"id": self.id, "name": self.name, "ast": self.ast.to_json()} + def to_json_bytes(self) -> bytes: + """`js_stringify_bytes(self.to_json())`, from the text the AST keeps.""" + id_, name = js_stringify_bytes(self.id), js_stringify_bytes(self.name) + return b'{"id":%b,"name":%b,"ast":%b}' % (id_, name, self.ast.to_json_bytes()) + class AppErroredQuery(RefusesNullOffWire): error: Literal["app"] = "app" diff --git a/src/zero_server/protocol/query_server.py b/src/zero_server/protocol/query_server.py index c65a1a9..3ecd550 100644 --- a/src/zero_server/protocol/query_server.py +++ b/src/zero_server/protocol/query_server.py @@ -6,9 +6,10 @@ from typing import Any, Final, Literal -from zero_server.protocol.custom_queries import TransformResponseBody +from zero_server.protocol.custom_queries import TransformedQuery, TransformResponseBody from zero_server.protocol.error import TransformFailedBody from zero_server.protocol.wire import dump_json, json_model +from zero_server.shared.js import js_stringify_bytes _UNSET: Final[Any] = object() @@ -39,5 +40,18 @@ def to_json(self) -> dict[str, Any]: out["userID"] = dump_json(self._user_id) return out + def to_json_bytes(self) -> bytes: + """`js_stringify_bytes(self.to_json())`, joined from each query's text rather than one walk of the whole.""" + queries = b",".join( + [ + q.to_json_bytes() if type(q) is TransformedQuery else js_stringify_bytes(q.to_json()) + for q in self.queries + ] + ) + out = b'{"kind":"QueryResponse","queries":[' + queries + b"]" + if self._user_id is not _UNSET: # python-only: a response built in Python without a user ID + out += b',"userID":' + js_stringify_bytes(dump_json(self._user_id)) + return out + b"}" + type QueryResponse = QuerySuccess | TransformFailedBody diff --git a/src/zero_server/query/builder.py b/src/zero_server/query/builder.py index 89e610d..48053f3 100644 --- a/src/zero_server/query/builder.py +++ b/src/zero_server/query/builder.py @@ -35,9 +35,11 @@ Correlation, Direction, LiteralReference, + NameMapper, SimpleCondition, System, insert_related, + map_ast, normalize_condition, ) from zero_server.query.expression import MISSING, ExpressionBuilder, _position, and_, cmp, simplify_condition @@ -86,6 +88,7 @@ class Query(Generic[R, C]): "_children", "_format", "_junction", + "_mapped", "_parent", "_schema", "_strict", @@ -113,6 +116,7 @@ def __init__( self._strict = strict self._parent: Query[Any, Any] | None = None self._children: dict[Any, weakref.ref[Any]] | None = None + self._mapped: tuple[NameMapper, AST] | None = None def __repr__(self) -> str: # python-only: for debugging return f"{type(self).__name__}({self._ast.to_json()!r})" @@ -152,6 +156,7 @@ def _copy[Q: Query[Any, Any]](self, cls: type[Q], ast: AST, format: Format) -> Q query._strict = self._strict query._parent = None query._children = None + query._mapped = None return query def _derive[T](self, key: Any, build: Callable[[], T]) -> T: @@ -263,17 +268,20 @@ def related[Q: Query[Any, Any]](self: Q, relationship: str, cb: SubqueryFn | Non if relationship.startswith(SUBQ_PREFIX): raise Exception(f'Relationship names may not start with "{SUBQ_PREFIX}". That is a reserved prefix.') related = self._relationship(relationship) + base_key = ("related_base", relationship) + q = self._find(base_key) if len(related) == 1: (hop,) = related - q = self._derive( - ("related_base", relationship), - lambda: self._new( - hop.dest_schema, - AST(table=hop.dest_schema, alias=relationship), - None, - _ONE if hop.cardinality == "one" else _MANY, - ), - ) + if q is None: + q = self._keep( + base_key, + self._new( + hop.dest_schema, + AST(table=hop.dest_schema, alias=relationship), + None, + _ONE if hop.cardinality == "one" else _MANY, + ), + ) sub = _as_query(cb(q) if cb else q) check(len(hop.source_field) >= 1, "The source of a relationship must specify at last 1 field") check(len(hop.dest_field) >= 1, "The destination of a relationship must specify at last 1 field") @@ -293,15 +301,16 @@ def related[Q: Query[Any, Any]](self: Q, relationship: str, cb: SubqueryFn | Non return self._with_relationship(relationship, sub, entry) first, second = related - q = self._derive( - ("related_base", relationship), - lambda: self._new( - second.dest_schema, - AST(table=second.dest_schema, alias=relationship), - relationship, - _ONE if second.cardinality == "one" else _MANY, - ), - ) + if q is None: + q = self._keep( + base_key, + self._new( + second.dest_schema, + AST(table=second.dest_schema, alias=relationship), + relationship, + _ONE if second.cardinality == "one" else _MANY, + ), + ) sub = _as_query(cb(q) if cb else q) _check_keys(first, second) found = self._find(("related", relationship, id(sub._ast))) @@ -374,16 +383,16 @@ def _exists( scalar: bool | None = None, ) -> Condition: related = self._relationship(relationship) + base_key = ("exists_base", relationship) + q = self._find(base_key) if len(related) == 1: (hop,) = related - check(len(hop.source_field) >= 1, "Invalid relationship") - check(len(hop.dest_field) >= 1, "Invalid relationship") - q = self._derive( - ("exists_base", relationship), - lambda: self._new( - hop.dest_schema, AST(table=hop.dest_schema, alias=f"{SUBQ_PREFIX}{relationship}"), None - ), - ) + check(hop.source_field and hop.dest_field, "Invalid relationship") + if q is None: + q = self._keep( + base_key, + self._new(hop.dest_schema, AST(table=hop.dest_schema, alias=f"{SUBQ_PREFIX}{relationship}"), None), + ) sub = _as_query(cb(q) if cb else q) return CorrelatedSubqueryCondition( related=CorrelatedSubquery( @@ -399,11 +408,11 @@ def _exists( first, second = related _check_keys(first, second) - alias = f"{SUBQ_PREFIX}zhidden_{relationship}" - q = self._derive( - ("exists_base", relationship), - lambda: self._new(second.dest_schema, AST(table=second.dest_schema, alias=alias), relationship), - ) + if q is None: + alias = f"{SUBQ_PREFIX}zhidden_{relationship}" + q = self._keep( + base_key, self._new(second.dest_schema, AST(table=second.dest_schema, alias=alias), relationship) + ) to_dest = _as_query(cb(q) if cb else q) return CorrelatedSubqueryCondition( related=CorrelatedSubquery( @@ -476,6 +485,12 @@ def _condition_key(condition: Condition) -> Any: keeps alive (`CorrelatedSubquery.source`).""" kind = type(condition) if kind is SimpleCondition: + left, right, op = condition.left, condition.right, condition.op # ty: ignore[unresolved-attribute] + # A column compared with a scalar, by far the usual condition, keyed flat: one tuple, not five. + if type(left) is ColumnReference and type(right) is LiteralReference and type(op) is str: + name, value = left.name, right.value + if type(name) is str and type(value) in _SCALAR_TYPES: + return name, op, type(value), value return ( kind, _position_key(condition.left), # ty: ignore[unresolved-attribute] @@ -484,10 +499,11 @@ def _condition_key(condition: Condition) -> Any: ) if kind is CorrelatedSubqueryCondition: related: CorrelatedSubquery = condition.related # ty: ignore[unresolved-attribute] + correlation = related.correlation return ( kind, - _tag(related.correlation.parent_field), - _tag(related.correlation.child_field), + _tag(correlation.parent_field), + _tag(correlation.child_field), _tag(related.hidden), _tag(related.system), _subquery_key(related), @@ -508,8 +524,7 @@ def _subquery_key(related: CorrelatedSubquery) -> Any: def _check_keys(first: Connection, second: Connection) -> None: - for key in (first.source_field, first.dest_field, second.source_field, second.dest_field): - check(len(key) >= 1, "Invalid relationship") + check(first.source_field and first.dest_field and second.source_field and second.dest_field, "Invalid relationship") def _is_number(value: object) -> bool: @@ -543,6 +558,16 @@ def as_query_internals(q: object) -> AnyQuery: return q +def mapped_ast(query: AnyQuery, mapper: NameMapper) -> AST: + """`map_ast(query.ast, mapper)`, kept on the query: an interned query comes back request after request.""" + mapped = query._mapped + if mapped is not None and mapped[0] is mapper: + return mapped[1] + ast = map_ast(query._ast, mapper) + query._mapped = (mapper, ast) + return ast + + def _as_query(q: object) -> AnyQuery: check(isinstance(q, Query), "Expected QueryImpl instance") assert isinstance(q, Query) diff --git a/src/zero_server/server/process_queries.py b/src/zero_server/server/process_queries.py index db5501a..932c3b6 100644 --- a/src/zero_server/server/process_queries.py +++ b/src/zero_server/server/process_queries.py @@ -13,7 +13,7 @@ from pydantic import TypeAdapter, ValidationError -from zero_server.protocol.ast import AST, map_ast +from zero_server.protocol.ast import AST from zero_server.protocol.custom_queries import ( AppErroredQuery, ParseErroredQuery, @@ -25,11 +25,11 @@ from zero_server.protocol.query_server import QuerySuccess from zero_server.protocol.schemas import refusal, transform_request_message_schema from zero_server.protocol.wire import parse_wire -from zero_server.query.builder import Query, as_query_internals +from zero_server.query.builder import Query, as_query_internals, mapped_ast from zero_server.query.registry import QueryParseError from zero_server.schema import Schema from zero_server.shared.error import error_name, get_error_details, get_error_message -from zero_server.shared.js import js_check_in, js_truthy +from zero_server.shared.js import js_check_in, js_stringify_bytes, js_truthy from zero_server.shared.json import JSONValue from zero_server.shared.json_parse import parse_body @@ -53,7 +53,14 @@ def server_ast(query: Query, schema: Schema) -> AST: Python only; upstream inlines `mapAST(ast, clientToServer(schema.tables))` in its query handler. Raises on a table or column `schema` does not have. """ - return map_ast(query.ast, schema.client_to_server_mapper) + return mapped_ast(query, schema.client_to_server_mapper) + + +def query_response_bytes(response: QueryResponse) -> bytes: + """`response` as the JSON text a TypeScript server sends: `JSON.stringify(response)`.""" + if isinstance(response, QuerySuccess): + return response.to_json_bytes() + return js_stringify_bytes(response.to_json()) async def handle_query_request( @@ -105,7 +112,7 @@ def transformed(request: TransformRequestQuery, result: object) -> Transformed: return errored(request, error) query = as_query_internals(result) try: - ast = map_ast(query.ast, mapper) + ast = mapped_ast(query, mapper) except Exception as error: logger.error("Failed to map AST", exc_info=error) raise diff --git a/src/zero_server/shared/js.py b/src/zero_server/shared/js.py index d69601f..dde95a2 100644 --- a/src/zero_server/shared/js.py +++ b/src/zero_server/shared/js.py @@ -170,14 +170,12 @@ def _may_differ(text: bytes) -> bool: outside = b"".join(parts[0::2]) if b"." in outside or outside.count(b"e") != outside.count(b"true") + outside.count(b"false"): return True - if not any(digit + b'":' in text for digit in _DIGITS): + # With the digits gone, a key of digits alone is an empty key: one scan rather than one per digit. + if b'"":' not in text.translate(None, b"0123456789"): return False return any(parts[i].isdigit() and parts[i + 1][:1] == b":" for i in range(1, len(parts) - 1, 2)) -_DIGITS = [str(d).encode() for d in range(10)] - - _SURROGATE = re.compile("[\ud800-\udfff]") diff --git a/tests/unit/test_protocol.py b/tests/unit/test_protocol.py index 070432f..d6c6156 100644 --- a/tests/unit/test_protocol.py +++ b/tests/unit/test_protocol.py @@ -50,6 +50,10 @@ def test_user_id_is_reported_only_when_set() -> None: assert QuerySuccess(queries=[], user_id=None).to_json() == {"kind": "QueryResponse", "queries": [], "userID": None} assert QuerySuccess(queries=[]).user_id is None assert QuerySuccess(queries=[], user_id="u").user_id == "u" + assert QuerySuccess(queries=[]).to_json_bytes() == b'{"kind":"QueryResponse","queries":[]}' + assert ( + QuerySuccess(queries=[], user_id=None).to_json_bytes() == b'{"kind":"QueryResponse","queries":[],"userID":null}' + ) @pytest.mark.parametrize("user", [{}, {"userID": None}, {"userID": "u1"}]) diff --git a/tests/wire/test_query.py b/tests/wire/test_query.py index ba4bdec..c743445 100644 --- a/tests/wire/test_query.py +++ b/tests/wire/test_query.py @@ -14,7 +14,7 @@ from tests.wire.programs import RETURNS, apply, make_handler, programs from zero_server import Query, create_builder, server_ast from zero_server.schema import Schema, load_schema -from zero_server.server.process_queries import handle_query_request +from zero_server.server.process_queries import handle_query_request, query_response_bytes from zero_server.shared.js import js_stringify @@ -64,6 +64,8 @@ def check(requests: list[Any], user_id: str | None) -> None: handle_query_request(handler=make_handler(schema), schema=schema, body=body, user_id=user_id) ) assert _comparable(actual.to_json()) == _comparable(expected) + # The text the HTTP endpoint sends, joined from each AST's own. + assert query_response_bytes(actual) == _comparable(expected).encode() check()