Native: query results as columns (arcadedb_n_exec_col / arcadedb_n_next_col), one crossing into a reused buffer (draft, stacked on #521)

bindings/python/src/arcadedb_embedded/_native/colres.py

@@ -0,0 +1,142 @@
+"""Columnar query results of the native library (``arcadedb_n_exec_col`` / ``arcadedb_n_next_col``): a frame of per-column buffers
+instead of JSON text. The layout is the one the columnar insert reads (int64 / float64 / bool bytes, UTF-8 strings as Arrow int32 offsets
+plus a byte blob, float32 vectors, bitmaps with bit i of byte i/8 least significant first); see ``ColResult.java`` for the frame. The rows
+come back as the same dicts the JSON reply gives, bit for bit, without parsing text.
+"""
+
+import json
+import struct
+from datetime import date, datetime, timedelta, timezone
+
+from .codec import RID
+
+MAGIC = 0x31524341  # b"ACR1"
+_HDR = struct.Struct(
+    "<IIQqII"
+)  # magic, JSON header bytes, frame bytes, cursor id, rows, flags
+_BASE = 32
+_ORDINAL_1970 = date(1970, 1, 1).toordinal()
+_EPOCH = datetime(1970, 1, 1)
+_EPOCH_UTC = datetime(1970, 1, 1, tzinfo=timezone.utc)
+_US = timedelta(microseconds=1)
+
+try:  # numpy is optional: it only speeds up big vector matrices
+    import numpy as _np
+except ImportError:  # pragma: no cover
+    _np = None
+_NUMPY_FROM = 4096  # floats in a matrix from which numpy's reshape().tolist() beats slicing in Python
+
+
+def is_frame(first_byte: int) -> bool:
+    """A reply starts with ``{`` (JSON) or with the frame magic (``A``)."""
+    return first_byte == 0x41
+
+
+def _strings(view, base, n, data, aux):
+    offs = view[base + data : base + data + 4 * (n + 1)].cast("i").tolist()
+    end = offs[n]
+    blob = bytes(view[base + aux : base + aux + end])
+    if blob.isascii():
+        text = blob.decode("ascii")
+        return [text[a:z] for a, z in zip(offs, offs[1:])]
+    return [blob[a:z].decode("utf-8") for a, z in zip(offs, offs[1:])]
+
+
+def _bitmap(view, base, off, n):
+    return bytes(view[base + off : base + off + ((n + 7) >> 3)])
+
+
+def _column(view, base, n, kind, dim, data, aux):
+    """The n cells of one column (a placeholder where the row has no value)."""
+    if kind == "q":
+        return view[base + data : base + data + 8 * n].cast("q").tolist()
+    if kind == "d":
+        return view[base + data : base + data + 8 * n].cast("d").tolist()
+    if kind == "s":
+        return _strings(view, base, n, data, aux)
+    if kind == "b":
+        return view[base + data : base + data + n].cast("?").tolist()
+    if kind == "f":
+        flat = view[base + data : base + data + 4 * n * dim]
+        if _np is not None and n * dim >= _NUMPY_FROM:
+            return _np.frombuffer(flat, dtype=_np.float32).reshape(n, dim).tolist()
+        flat = flat.cast("f").tolist()
+        return [flat[i : i + dim] for i in range(0, n * dim, dim)]
+    if kind == "v":
+        offs = view[base + data : base + data + 4 * (n + 1)].cast("i").tolist()
+        flat = view[base + aux : base + aux + 4 * offs[n]].cast("f").tolist()
+        return [flat[a:z] for a, z in zip(offs, offs[1:])]
+    if kind == "R":
+        return [RID(t) for t in _strings(view, base, n, data, aux)]
+    if kind == "D":
+        ords = view[base + data : base + data + 8 * n].cast("q").tolist()
+        return [date.fromordinal(x + _ORDINAL_1970) for x in ords]
+    if kind == "T":
+        us = view[base + data : base + data + 8 * n].cast("q").tolist()
+        return [_EPOCH + x * _US for x in us]
+    if kind == "Z":
+        us = view[base + data : base + data + 8 * n].cast("q").tolist()
+        return [_EPOCH_UTC + x * _US for x in us]
+    if kind == "n":
+        return [None] * n
+    raise ValueError(f"unknown column kind {kind!r} in a columnar reply")
+
+
+def decode(view) -> dict:
+    """The reply dict (``ok``, ``rs``, ``rows``, ``meta``, ``drained``, maybe ``rerr``) of a frame; ``view`` is a byte memoryview of it.
+    Nothing of ``view`` is kept: the buffer it points into is reused by the next call on the thread.
+    """
+    magic, hlen, total, rs, n, flags = _HDR.unpack_from(view, 0)
+    if magic != MAGIC:
+        raise ValueError("not a columnar reply")
+    head = json.loads(bytes(view[_BASE : _BASE + hlen]))
+    base = _BASE + hlen
+    names, cols, sparse = [], [], []
+    for name, kind, dim, data, aux, valid, present in head["c"]:
+        vals = _column(view, base, n, kind, dim, data, aux)
+        if valid:
+            vb = _bitmap(view, base, valid, n)
+            vals = [
+                v if (vb[i >> 3] >> (i & 7)) & 1 else None for i, v in enumerate(vals)
+            ]
+        names.append(name)
+        cols.append(vals)
+        if present:
+            sparse.append((name, vals, _bitmap(view, base, present, n)))
+    orders = {r: o for r, o in head.get("o", ())}
+    if not sparse and not orders:
+        rows = (
+            [dict(zip(names, r)) for r in zip(*cols)]
+            if cols
+            else [{} for _ in range(n)]
+        )
+    else:
+        # some rows lack a property or list theirs in another order (record rows): build each row from the columns whose bit says the property exists
+        present = {name: pb for name, _, pb in sparse}
+        keep = [(name, vals, present.get(name)) for name, vals in zip(names, cols)]
+        rows = []
+        for i in range(n):
+            row = {}
+            o = orders.get(i)
+            for name, vals, pb in keep if o is None else [keep[j] for j in o]:
+                if pb is None or (pb[i >> 3] >> (i & 7)) & 1:
+                    row[name] = vals[i]
+            rows.append(row)
+    m = head["m"]
+    if not m:
+        meta = [0] * n
+    else:
+        kinds = bytes(view[base + m[0] : base + m[0] + n])
+        rids = _strings(view, base, n, m[1], m[2])
+        types = _strings(view, base, n, m[3], m[4])
+        meta = [[r, t, chr(k)] if k else 0 for r, t, k in zip(rids, types, kinds)]
+    reply = {
+        "ok": True,
+        "rs": rs,
+        "rows": rows,
+        "meta": meta,
+        "drained": bool(flags & 1),
+    }
+    if "rerr" in head:
+        reply["rerr"] = head["rerr"]
+    return reply

bindings/python/src/arcadedb_embedded/_native/database.py

@@ -199,7 +199,14 @@ def _fill(self) -> bool:
         if not self._more:
             return False
         try:
-            reply = loads(lib.call("arcadedb_n_next", self._rs, _NEXT_BATCH))
+            if lib.COLRES and _BINARY_REPLY:
+                reply = lib.next_col(self._rs, _NEXT_BATCH)
+            elif lib.NEXT_BIN and _BINARY_REPLY:
+                reply = loads(
+                    *lib.call_bin("arcadedb_n_next_bin", self._rs, _NEXT_BATCH)
+                )
+            else:
+                reply = loads(lib.call("arcadedb_n_next", self._rs, _NEXT_BATCH))
         except ArcadeDBError:
             raise
         if not reply.get("ok"):
@@ -468,6 +475,23 @@ def _exec(
         self, flags: int, prefix: str, language: str, command: str, args
     ) -> NativeResultSet:
         self._check_not_closed()
+        if not flags & 1 and lib.COLRES and _BINARY_REPLY:
+            # N-R2: the rows come back as columns (JSON for a batch of fewer than lib.COLRES_MIN rows or one that columns cannot carry), in the
+            # thread's reusable reply buffer; the vectors of a JSON batch still travel in the float32 tail (flag 4)
+            text, buf, count = encode_params(args, binary=True)
+            reply = lib.exec_col(
+                self._h,
+                flags | 4,
+                lib.bc(language),
+                lib.bc(command),
+                lib.b(text),
+                buf,
+                count,
+                _FIRST_BATCH,
+            )
+            if not reply.get("ok"):
+                raise _error(prefix, reply)
+            return NativeResultSet(self, reply)
         if flags & 1 and _LAZY_COMMAND_ROWS and lib.CMD_PATH:
             # phase 8: the cursor id comes back as a long (no reply text to free and decode); the rows are read late or never
             text, buf, count = encode_params(args, binary=True)

bindings/python/src/arcadedb_embedded/_native/lib.py

@@ -8,6 +8,8 @@
 from pathlib import Path
 
 from ..exceptions import ArcadeDBError
+from . import colres
+from .codec import loads
 
 _LIB_NAMES = {
     "linux": "libarcadedb.so",
@@ -28,6 +30,18 @@
 INSERT_ROWS = os.environ.get("ARCADEDB_NATIVE_INSERT_ROWS", "1") != "0"
 # ARCADEDB_NATIVE_EXEC_MANY: the rows of an insert_many chunk that JSON text cannot carry unchanged go in one crossing (arcadedb_n_exec_many), not one call each
 EXEC_MANY = os.environ.get("ARCADEDB_NATIVE_EXEC_MANY", "1") != "0"
+# ARCADEDB_NATIVE_NEXT_BIN: the batches after the first read through arcadedb_n_next_bin, float arrays of 8 or more elements in a float32 tail (as flag 4 of arcadedb_n_exec does for the first batch), not as number text
+NEXT_BIN = os.environ.get("ARCADEDB_NATIVE_NEXT_BIN", "1") != "0"
+# ARCADEDB_NATIVE_COLRES: query results come back as columns (arcadedb_n_exec_col / arcadedb_n_next_col: one frame of per-column buffers, in a
+# buffer of the caller's that is reused, no arcadedb_free) when the library has them; 0 forces the JSON text of the paths above, in the same image.
+# ARCADEDB_NATIVE_COLRES_MIN: batches of fewer rows answer as JSON text (a point lookup gains nothing from columns); 1 = columns for every batch.
+COLRES = os.environ.get("ARCADEDB_NATIVE_COLRES", "1") != "0"
+COLRES_MIN = int(os.environ.get("ARCADEDB_NATIVE_COLRES_MIN", "16"))
+col_frames = 0  # replies that came back as columnar frames / as JSON through the columnar entry points (counters for tests and probes)
+col_json = 0
+_REPLY_CAP = (
+    1 << 19
+)  # per-thread reply buffer; a bigger reply comes in a block the image allocates and Python copies and frees
 
 _lock = threading.Lock()
 _lib = None
@@ -159,6 +173,9 @@ class _Params(ctypes.Structure):
             "arcadedb_n_exec_bin": [P, L, I, S, S, S, I, S, L],
             "arcadedb_n_insert_cols": [P, L, S, S, S, L],
             "arcadedb_n_next": [P, L, I],
+            "arcadedb_n_next_bin": [P, L, I],
+            "arcadedb_n_exec_col": [P, L, I, S, S, S, I, S, L, P, L, L],
+            "arcadedb_n_next_col": [P, L, I, P, L, L],
             "arcadedb_n_rs_close": [P, L],
             "arcadedb_n_cmd": [P, L, I, S, S, S, S, L, I],
             "arcadedb_n_last_error": [P],
@@ -195,7 +212,7 @@ class _Params(ctypes.Structure):
                 )
                 else ctypes.c_longlong if name == "arcadedb_n_cmd" else ctypes.c_void_p
             )
-        global BINARY_PARAMS, CMD_PATH, CLOSE_BATCH, INSERT_ROWS, EXEC_MANY
+        global BINARY_PARAMS, CMD_PATH, CLOSE_BATCH, INSERT_ROWS, EXEC_MANY, NEXT_BIN, COLRES
         if not hasattr(
             lib, "arcadedb_n_exec_bin"
         ):  # a library built before the binary path
@@ -208,6 +225,14 @@ class _Params(ctypes.Structure):
             INSERT_ROWS = False
         if not hasattr(lib, "arcadedb_n_exec_many"):
             EXEC_MANY = False
+        if not hasattr(
+            lib, "arcadedb_n_next_bin"
+        ):  # a library built before the binary tail of later batches
+            NEXT_BIN = False
+        if not hasattr(lib, "arcadedb_n_exec_col") or not hasattr(
+            lib, "arcadedb_n_next_col"
+        ):  # a library built before columnar results
+            COLRES = False
         _isolate = isolate
         _creator = threading.get_ident()
         _local.thread = _Thread(thread, owned=False)
@@ -268,6 +293,98 @@ def call_bin(name: str, *args):
     return raw.decode("utf-8"), floats
 
 
+def _reply_buffer():
+    r = getattr(_local, "reply", None)
+    if r is None:
+        ba = bytearray(_REPLY_CAP)
+        r = _local.reply = (
+            ba,
+            ctypes.addressof(ctypes.c_char.from_buffer(ba)),
+            memoryview(ba),
+        )
+    return r
+
+
+def _json_tail(raw, floats_at):
+    """``(text, floats)`` of a JSON reply (``raw``: its bytes up to the NUL); ``floats_at(offset, count)`` reads the float32 tail behind the NUL."""
+    floats = None
+    if raw.endswith(b"}"):
+        i = raw.rfind(b',"bin":', max(0, len(raw) - 40))
+        if i >= 0:
+            import numpy as np
+
+            n = int(raw[i + 7 : -1])
+            floats = np.frombuffer(
+                floats_at((len(raw) + 1 + 3) & ~3, n), dtype=np.float32
+            )
+            raw = raw[:i] + b"}"
+    return raw.decode("utf-8"), floats
+
+
+def _col_reply(lib, th, ptr, addr, view):
+    """The decoded reply behind ``ptr``: a columnar frame or the JSON text (with the float32 tail) of ``call_bin``. ``ptr == addr`` is the
+    thread's own reply buffer (valid until the next call on the thread, nothing to free); anything else is a block the image allocated
+    (an oversize reply, or an error): copied out, then freed."""
+    global col_frames, col_json
+    if ptr == addr:
+        if view[0] == 0x41:
+            col_frames += 1
+            return colres.decode(view[: int.from_bytes(view[8:16], "little")])
+        col_json += 1
+        return loads(
+            *_json_tail(ctypes.string_at(ptr), lambda off, n: view[off : off + 4 * n])
+        )
+    frame = None
+    try:
+        if ctypes.string_at(ptr, 1) == b"A":
+            frame = ctypes.string_at(
+                ptr, int.from_bytes(ctypes.string_at(ptr + 8, 8), "little")
+            )
+        else:
+            text, floats = _json_tail(
+                ctypes.string_at(ptr), lambda off, n: ctypes.string_at(ptr + off, 4 * n)
+            )
+    finally:
+        lib.arcadedb_free(th, ptr)
+    if frame is not None:
+        col_frames += 1
+        return colres.decode(memoryview(frame))
+    col_json += 1
+    return loads(text, floats)
+
+
+def exec_col(handle, flags, language, command, params, buf, count, first) -> dict:
+    """``arcadedb_n_exec_col``: the decoded reply of the first batch (columns, or JSON for a small or too mixed batch). One crossing, the reply in
+    the thread's reusable buffer: no ``arcadedb_free`` for a reply that fits it."""
+    lib = _lib or _load()
+    th = _thread_handle()
+    _, addr, view = _reply_buffer()
+    ptr = lib.arcadedb_n_exec_col(
+        th,
+        handle,
+        flags,
+        language,
+        command,
+        params,
+        first,
+        buf,
+        count,
+        addr,
+        _REPLY_CAP,
+        COLRES_MIN,
+    )
+    return _col_reply(lib, th, ptr, addr, view)
+
+
+def next_col(rs_id, count) -> dict:
+    """``arcadedb_n_next_col``: the decoded next batch of a cursor."""
+    lib = _lib or _load()
+    th = _thread_handle()
+    _, addr, view = _reply_buffer()
+    ptr = lib.arcadedb_n_next_col(th, rs_id, count, addr, _REPLY_CAP, COLRES_MIN)
+    return _col_reply(lib, th, ptr, addr, view)
+
+
 def call_void(name: str, *args) -> None:
     """Run an entry point that returns nothing (``arcadedb_n_set_upcall``)."""
     lib = _lib or _load()