perf: send query responses up to 3x faster by reusing mapped ASTs and their JSON

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

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"]))
 

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.

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"

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

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)

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

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]")