diff --git a/bindings/python/src/arcadedb_embedded/_native/colres.py b/bindings/python/src/arcadedb_embedded/_native/colres.py new file mode 100644 index 00000000000..972409cb040 --- /dev/null +++ b/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( + " 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 diff --git a/bindings/python/src/arcadedb_embedded/_native/database.py b/bindings/python/src/arcadedb_embedded/_native/database.py index 68fe220b894..4519f7c3ac9 100644 --- a/bindings/python/src/arcadedb_embedded/_native/database.py +++ b/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) diff --git a/bindings/python/src/arcadedb_embedded/_native/lib.py b/bindings/python/src/arcadedb_embedded/_native/lib.py index 8ac33145143..956a4ac2d76 100644 --- a/bindings/python/src/arcadedb_embedded/_native/lib.py +++ b/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() diff --git a/bindings/python/tests/test_native_columnar_results.py b/bindings/python/tests/test_native_columnar_results.py new file mode 100644 index 00000000000..7d84f91ac76 --- /dev/null +++ b/bindings/python/tests/test_native_columnar_results.py @@ -0,0 +1,350 @@ +"""Query results as columns (N-R2): the rows a result gives back are the same values, types and key order whether the library answers with JSON text +or with a columnar frame, for mixed types with nulls, unicode, integers near 2**31 and 2**53, doubles with -0.0, inf and nan, float32 vectors of +every length, record ids, dates and datetimes, in one batch and in several, for point lookups, for records with missing properties and for results +the columns cannot carry (nested maps and lists, a column that mixes int and double). On the native library the same query runs three ways in one +process: JSON (ARCADEDB_NATIVE_COLRES=0), columns for every batch (COLRES_MIN=1) and the default; the JPype backend only checks the expected values. +""" + +import math +import os +import struct +from contextlib import contextmanager +from datetime import date, datetime + +NATIVE = os.environ.get("ARCADEDB_BACKEND", "jpype").lower() == "native" +if NATIVE: + from arcadedb_embedded._native import lib +else: + lib = None + +TEXTS = [ + "", + "plain", + "café 日本語", + "\U0001f600 emoji \U0001f468‍\U0001f469", + 'quote " backslash \\ nl \n tab \t ctrl \u0001 slash /', + "x" * 300, + None, +] +LONGS = [ + -(1 << 63), + -1, + 0, + (1 << 31) - 1, + 1 << 31, + 1 << 53, + (1 << 53) + 1, + (1 << 62), + (1 << 63) - 1, +] +DOUBLES = [ + -0.0, + 0.0, + float("inf"), + float("-inf"), + float("nan"), + 1e-300, + 0.1, + 1.7976931348623157e308, + 5e-324, + 123456789.12345679, +] + + +def f32(x): + return struct.unpack("f", struct.pack("f", x))[0] + + +def bits(x): + return "nan" if x != x else struct.pack(" f0 + ) # carried as a frame (the rows are all of one kind per column) + + +def test_record_rows_keep_their_meta_and_property_order(temp_db): + db, rows = _db(temp_db, 30) + got = _modes(db, "SELECT FROM T ORDER BY id") + assert [r["id"] for r in got] == list(range(30)) + with mode(True, 1): + rs = db.query("sql", "SELECT FROM T ORDER BY id") + first = next(iter(rs)) + assert first.get_rid() is not None and first.get_element() is not None + + +def test_record_ids_and_types_as_values(temp_db): + db, _ = _db(temp_db, 12) + got = _modes(db, "SELECT @rid AS r, @type AS t, id FROM T ORDER BY id") + assert all(str(r["r"]).startswith("#") and r["t"] == "T" for r in got) + assert all(type(r["r"]).__name__ == "RID" for r in got) + + +def test_batches_of_a_big_result_are_the_same_as_one(temp_db): + db, rows = _db(temp_db, 5000) # 256 in the first reply, then 2048 + 2048 + 648 + got = _modes(db, "SELECT id, big, d, s, v FROM T ORDER BY id") + assert len(got) == 5000 + assert [r["id"] for r in got] == list(range(5000)) + for want, have in zip(rows, got): + assert ( + have["big"] == want["big"] + and bits(have["d"]) == bits(want["d"]) + and have.get("s") == want.get("s") + ) + assert (list(have["v"]) if have.get("v") is not None else None) == want.get("v") + + +def test_point_lookups_and_empty_results(temp_db): + db, rows = _db(temp_db, 50) + got = _modes(db, "SELECT id, s, d FROM T WHERE id = :id", {"id": 3}) + assert ( + len(got) == 1 and got[0]["id"] == 3 and bits(got[0]["d"]) == bits(rows[3]["d"]) + ) + assert _modes(db, "SELECT FROM T WHERE id < 0") == [] + assert _modes(db, "SELECT 1 AS one, 'a' AS two") == [{"one": 1, "two": "a"}] + agg = _modes(db, "SELECT count(*) AS c, avg(id) AS a FROM T") + assert agg[0]["c"] == 50 + + +def test_records_with_missing_and_null_properties(temp_db): + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE Q") + for i in range(30): + row = {"a": i} + if i % 3: + row["b"] = f"b{i}" + if i % 4: + row["c"] = i * 0.5 + if i % 5 == 0: + row["extra"] = [i, "x"] # a list: the batch it is in goes as JSON + db.insert_many("Q", [row]) + got = _modes(db, "SELECT FROM Q ORDER BY a") + assert [r["a"] for r in got] == list(range(30)) + assert all(("b" in r) == bool(r["a"] % 3) for r in got) + proj = _modes(db, "SELECT a, b, c FROM Q ORDER BY a") + assert all( + set(r) == {"a", "b", "c"} for r in proj + ) # a projection names every column, null where the property is missing + + +def test_results_the_columns_cannot_carry_still_come_back(temp_db): + db, _ = _db(temp_db, 20) + f0, j0 = _counts() + nested = _modes( + db, "SELECT id, [id, id + 1] AS pair, {'k': id} AS m FROM T ORDER BY id" + ) + assert nested[3]["pair"] == [3, 4] and nested[3]["m"] == {"k": 3} + mixed = _modes(db, "SELECT id, ifnull(NULL, d) AS z FROM T ORDER BY id") # doubles + assert len(mixed) == 20 + with mode(True, 1): + _run(db, "SELECT id, [id, id + 1] AS pair FROM T ORDER BY id") + if lib is not None and hasattr(lib, "col_json"): + assert _counts()[1] > j0 # answered as JSON through the columnar entry points + + +def test_a_cursor_read_row_by_row_matches(temp_db): + db, rows = _db(temp_db, 600) + for label, (c, m) in {"json": (False, None), "cols": (True, 1)}.items(): + with mode(c, m): + seen = [ + r.get("id") for r in db.query("sql", "SELECT id, s FROM T ORDER BY id") + ] + assert seen == list(range(600)), label + + +def test_record_rows_of_several_batches(temp_db): + """Record rows (a meta table next to the columns) in batches of 256 and 2048 rows: the first N-R2 build wrote bitmaps for the meta columns + and failed from 65 rows on (found by the transport probe's SELECT FROM D, not by the 20-row cases). + """ + db, rows = _db(temp_db, 5000) + got = _modes(db, "SELECT FROM T ORDER BY id") + assert [r["id"] for r in got] == list(range(5000)) + with mode(True, 1): + rs = db.query("sql", "SELECT FROM T ORDER BY id") + first, count = next(iter(rs)), 1 + sum(1 for _ in rs) + assert count == 5000 and first.get_rid() is not None + + +def test_the_probes_wide_table_in_every_mode(temp_db): + """The transport probe's table and queries (id, grp, val, cat, txt, vec of 32 floats; 3,000 rows), record rows and a projection.""" + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE D") + for p, t in ( + ("id", "INTEGER"), + ("grp", "INTEGER"), + ("val", "DOUBLE"), + ("cat", "STRING"), + ("txt", "STRING"), + ("vec", "ARRAY_OF_FLOATS"), + ): + db.command("sql", f"CREATE PROPERTY D.{p} {t}") + db.insert_many( + "D", + [ + { + "id": i, + "grp": i % 10, + "val": i * 0.37, + "cat": f"category-{i % 20}", + "txt": "t" + "x" * (5 + i % 25), + "vec": [f32((i + j) / 7.0) for j in range(32)], + } + for i in range(3000) + ], + ) + for sql in ( + "SELECT FROM D LIMIT 3000", + "SELECT id, grp, val, cat, txt, vec FROM D LIMIT 3000", + ): + got = _modes(db, sql) + assert len(got) == 3000 + assert all(len(r["vec"]) == 32 for r in got) diff --git a/bindings/python/tests/test_native_result_transport.py b/bindings/python/tests/test_native_result_transport.py new file mode 100644 index 00000000000..12edbfcb33d --- /dev/null +++ b/bindings/python/tests/test_native_result_transport.py @@ -0,0 +1,173 @@ +"""Query results across the native boundary: the reply text (and the float32 tail of a reply with vectors) the image hands to Python +is copied out of the image in bulk, not one byte at a time. Every value has to come back exactly as it went in, for replies of +every size: tiny, several batches (arcadedb_n_exec returns the first 256 rows, arcadedb_n_next the following 2048 at a time) and +single rows of about a megabyte of multi-byte text. The tests hold on both backends; only the native one has this copy. +""" + +import struct + +DIM = 16 # 8 or more floats per vector: they travel in the float32 tail of the first batch, as JSON numbers in the later ones + +TEXTS = [ + "", + "plain", + "café 日本語 \U0001f600 end", # 2-, 3- and 4-byte UTF-8 sequences + 'quote " backslash \\ slash / newline \n tab \t nul-free \u0001 control', + "x" * 300, +] + + +def f32(x): + """The double a float32 value widens to (what the engine returns for an ARRAY_OF_FLOATS element).""" + return struct.unpack("f", struct.pack("f", x))[0] + + +def _rows(n): + rows = [] + for i in range(n): + rows.append( + { + "id": i, + "big": (1 << 62) - i if i % 2 else -(1 << 62) + i, + "val": (i + 0.5) / 7.0, + "txt": TEXTS[i % len(TEXTS)], + "vec": [f32((i * DIM + k) / 1000.0 - 3.0) for k in range(DIM)], + } + ) + return rows + + +def _typed_db(temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE R") + for p, t in ( + ("id", "INTEGER"), + ("big", "LONG"), + ("val", "DOUBLE"), + ("txt", "STRING"), + ("vec", "ARRAY_OF_FLOATS"), + ): + temp_db.command("sql", f"CREATE PROPERTY R.{p} {t}") + return temp_db + + +def test_a_reply_of_several_batches_round_trips_exactly(temp_db): + db = _typed_db(temp_db) + rows = _rows( + 5000 + ) # 256 in the first reply (vectors in the float tail), then 2048 + 2048 + 648 (vectors as numbers) + db.insert_many("R", rows) + got = db.query("sql", "SELECT id, big, val, txt, vec FROM R ORDER BY id").to_list() + assert len(got) == len(rows) + for want, have in zip(rows, got): + assert have["id"] == want["id"] + assert have["big"] == want["big"] + assert have["val"] == want["val"] + assert have["txt"] == want["txt"] + assert list(have["vec"]) == want["vec"] + + +def test_a_small_reply_with_a_vector_round_trips_exactly(temp_db): + db = _typed_db(temp_db) + rows = _rows(3) + db.insert_many("R", rows) + got = db.query("sql", "SELECT id, vec, txt FROM R ORDER BY id").to_list() + assert [list(r["vec"]) for r in got] == [r["vec"] for r in rows] + assert [r["txt"] for r in got] == [r["txt"] for r in rows] + + +def test_rows_of_about_a_megabyte_of_multibyte_text_round_trip(temp_db): + db = _typed_db(temp_db) + unit = "aé日\U0001f600" # 1 + 2 + 3 + 4 bytes + rows = [ + {"id": i, "txt": unit * (250_000 + i)} for i in range(4) + ] # ~1.2 MB of UTF-8 per row + db.insert_many("R", rows) + got = db.query("sql", "SELECT id, txt FROM R ORDER BY id").to_list() + assert [(r["id"], r["txt"]) for r in got] == [(r["id"], r["txt"]) for r in rows] + + +def test_empty_and_one_row_replies(temp_db): + db = _typed_db(temp_db) + assert db.query("sql", "SELECT FROM R").to_list() == [] + assert db.query("sql", "SELECT '' AS s").to_list() == [{"s": ""}] + assert db.query("sql", "SELECT 1 AS x").first().get("x") == 1 + + +# ---- vectors on every batch (first reply and arcadedb_n_next batches), bit-exact, with short, empty and missing vectors + +SPECIAL = [ + float("inf"), + float("-inf"), + -0.0, + 0.0, + 1e-45, + 3.4028235e38, + 0.1, + -123456.789, +] # 8 values: a float32 tail on the binary path + + +def _bits(x): + """The 8 bytes of a double as hex (any NaN is 'nan'): bit-exact comparison, so -0.0 differs from 0.0.""" + return "nan" if x != x else struct.pack(" batchOf(rsId, cursor, max)); + } catch (final Throwable t) { + RESULT_SETS.remove(rsId); + closeQuietly(rs); + return err(t); + } + } + @CEntryPoint(name = "arcadedb_n_rs_close") public static CCharPointer rsClose(final IsolateThread thread, final long rsId) { closeQuietly(RESULT_SETS.remove(rsId)); @@ -487,11 +583,19 @@ private static CCharPointer replyOf(final ResultSet rs, final int max, final int } if ((flags & 4) == 0) return out(batch(rs, max)); + return withFloatTail(() -> batch(rs, max)); + } + + /** + * Runs encode with the float32 sink set (float[] values of 8 or more elements go to the tail of the reply as {"\u0001":["B",off,n]}) and returns + * the reply: out(text) when no vector was met, else the text with ,"bin":total, a NUL, padding to 4 bytes and the floats. + */ + private static CCharPointer withFloatTail(final java.util.function.Supplier encode) { final NativeCodec.FloatSink sink = new NativeCodec.FloatSink(); NativeCodec.SINK.set(sink); final String reply; try { - reply = batch(rs, max); + reply = encode.get(); } finally { NativeCodec.SINK.remove(); } @@ -500,13 +604,12 @@ private static CCharPointer replyOf(final ResultSet rs, final int max, final int final byte[] b = (reply.substring(0, reply.length() - 1) + ",\"bin\":" + sink.n + "}").getBytes(StandardCharsets.UTF_8); final int off = (b.length + 1 + 3) & ~3; final CCharPointer p = UnmanagedMemory.malloc(off + 4 * sink.n); - for (int i = 0; i < b.length; i++) - p.write(i, b[i]); - for (int i = b.length; i < off; i++) - p.write(i, (byte) 0); - final CFloatPointer fp = (CFloatPointer) ((org.graalvm.word.Pointer) p).add(off); - for (int i = 0; i < sink.n; i++) - fp.write(i, sink.a[i]); + // bulk copies (Unsafe.copyMemory), not one p.write per byte: the text, the zero padding up to the 4-byte boundary, then the float32 tail in native order + final java.nio.ByteBuffer bb = CTypeConversion.asByteBuffer(p, (int) (off + 4L * sink.n)); + bb.put(b, 0, b.length); + while (bb.position() < off) + bb.put((byte) 0); + bb.order(java.nio.ByteOrder.nativeOrder()).asFloatBuffer().put(sink.a, 0, sink.n); return p; } @@ -683,9 +786,8 @@ static Chars chars(final CCharPointer p) { static CCharPointer out(final String s) { final byte[] b = s.getBytes(StandardCharsets.UTF_8); final CCharPointer p = UnmanagedMemory.malloc(b.length + 1); - for (int i = 0; i < b.length; i++) - p.write(i, b[i]); - p.write(b.length, (byte) 0); + // one bulk copy (Unsafe.copyMemory) and the terminating zero; this was a Java loop of p.write(i, b[i]), one byte per iteration + CTypeConversion.asByteBuffer(p, b.length + 1).put(b, 0, b.length).put((byte) 0); return p; } } diff --git a/native-spike/src/main/java/org/humemai/nativelib/ColResult.java b/native-spike/src/main/java/org/humemai/nativelib/ColResult.java new file mode 100644 index 00000000000..67a7437be90 --- /dev/null +++ b/native-spike/src/main/java/org/humemai/nativelib/ColResult.java @@ -0,0 +1,718 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Phase 9 (N-R2): query results as columns. The read-side counterpart of the columnar insert (BulkOps): a batch of rows is laid out as + * one frame, a fixed header, a small JSON header with the column names and kinds, then one buffer per column in the layout BulkOps reads + * (native-endian int64 / float64 / bool bytes, UTF-8 strings as Arrow int32 offsets plus a byte blob, float32 vector matrix, bitmaps with + * bit i of byte i/8, least significant first). Python restores the rows from the buffers (_native/colres.py) without parsing text. + * + * Frame (all numbers native-endian, every buffer starts on an 8-byte boundary): + * 0 u32 magic "ACR1" 4 u32 length of the JSON header (padded with spaces to a multiple of 8) 8 u64 frame bytes + * 16 i64 cursor id (0 = drained) 24 u32 rows 28 u32 flags (bit 0: drained) 32 JSON header then the buffers + * JSON header: {"c":[[name,kind,dim,data,aux,valid,present],..],"m":0|[kind,rid,ridAux,type,typeAux][,"o":[[row,[column,..]],..]][,"rerr":{..}]}. Offsets are relative to + * the first byte after the header and 0 means "none" (the buffers start at 8). Kinds: q int64, d float64, b bool, D days (int64), T and Z + * microseconds since the epoch (int64; T a wall clock, Z UTC), s string, R record id as text, v float[] with int32 offsets, f float[] of + * one length (dim floats per row), n only nulls. valid: bit set = the row has a value; present: bit set = the property exists in the row + * (valid without present = null, neither = the row has no such property); no bitmap = all rows. m: the record rows' meta, kind = one byte per + * row (0 none, v, e, d), rid and type as string columns. o: rows whose properties come in another order than the columns (the engine does not + * order the properties of schemaless records alike), as the column indexes in the row's order; a batch with more of them than half its rows goes as JSON. + * + * A batch that cannot be carried (a nested map, list or record, a value type the buffers do not have, a column that mixes kinds, property + * order that differs between rows, more than MAX_COLS columns, a mostly empty table, an unpaired surrogate) is answered as the JSON text of + * arcadedb_n_next_bin, so every result still comes back. So is a batch of fewer than minRows rows: a point lookup gains nothing from columns. + */ +package org.humemai.nativelib; + +import com.arcadedb.database.Document; +import com.arcadedb.database.Identifiable; +import com.arcadedb.graph.Edge; +import com.arcadedb.graph.Vertex; +import com.arcadedb.query.sql.executor.Result; +import com.arcadedb.query.sql.executor.ResultSet; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Calendar; +import java.util.Date; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +final class ColResult { + static final int MAGIC = 0x31524341; // 'A' 'C' 'R' '1' as a little-endian int + static final int HEADER = 32; + static final int MAX_COLS = 256; + static final long MAX_BYTES = 4L << 20; // column bytes per batch, as NativeCodec.MAX_BATCH_CHARS for the text + private static final byte[] ZEROS = new byte[8]; + /** Why the last batch of this process went as JSON (tests and the JVM gate read it; racy by design). */ + static String lastFallback; + + private ColResult() { + } + + /** The batch cannot be carried as columns; the rows go as JSON. */ + private static final class Unsupported extends RuntimeException { + Unsupported(final String why) { + super(why, null, false, false); + } + } + + /** One reply: a JSON batch (bytes) or a columnar frame (the builder). */ + static final class Reply { + boolean drained; + int rows; + byte[] json; // text, or text + NUL + padding + float32 tail, without the terminating zero + Frame frame; + + int size() { + return json != null ? json.length + 1 : frame.total; + } + + void writeTo(final ByteBuffer bb) { + if (json != null) { + bb.put(0, json, 0, json.length); + bb.put(json.length, (byte) 0); + } else + frame.write(bb); + } + } + + /** + * Reads up to max rows of rs and answers them: a columnar frame when the rows are carried and there are at least minRows of them, else the + * JSON batch (with the float32 tail when tail is set). id is the cursor id shown to the caller (0 once drained). + */ + static Reply read(final long id, final ResultSet rs, final int max, final int minRows, final boolean tail) { + final NativeCodec.FloatSink sink = tail ? new NativeCodec.FloatSink() : null; + if (sink != null) + NativeCodec.SINK.set(sink); + try { + final List rows = new ArrayList<>(); + Builder cb = null; + boolean colOk = minRows >= 0; + final int start = Math.min(Math.max(1, minRows), Math.max(1, max)); // rows in hand before the columns are started + StringBuilder sb = null, meta = null; // JSON mode, once the columns are given up + int encoded = 0; + boolean drained = false; + String readError = null; + try { + while (rows.size() < max) { + if (!rs.hasNext()) { + drained = true; + break; + } + final Result row = rs.next(); + rows.add(row); + if (colOk && rows.size() >= start) { + try { + if (cb == null) { + cb = new Builder(); + for (final Result r : rows) + cb.add(r); + } else + cb.add(row); + } catch (final Throwable t) { // Unsupported, or a value the engine would not give up: the JSON path meets it again and reports it + lastFallback = t.toString(); + colOk = false; + cb = null; + } + } + if (cb != null && cb.bytes >= MAX_BYTES) + break; + if (!colOk) { + if (sb == null) { + sb = new StringBuilder(); + meta = new StringBuilder(); + } + // JSON mode: encode what has not been encoded yet (all rows so far on the first turn), then keep the 4M-char cap of the text path + while (encoded < rows.size()) { + final int rowStart = sb.length(), metaStart = meta.length(); + try { + if (encoded > 0) { + sb.append(','); + meta.append(','); + } + NativeCodec.appendRow(sb, meta, rows.get(encoded)); + } catch (final Throwable t) { + sb.setLength(rowStart); + meta.setLength(metaStart); + throw t; + } + encoded++; + } + if (sb.length() >= NativeCodec.MAX_BATCH_CHARS) + break; + } + } + } catch (final Throwable t) { + readError = ArcadeNative.errorJson(t); + drained = true; + } + final Reply reply = new Reply(); + reply.drained = drained; + reply.rows = rows.size(); + final long shown = drained ? 0 : id; + if (cb != null && colOk && !cb.sparse()) { // a builder that is alive holds every pulled row + reply.rows = cb.n; + reply.frame = cb.finish(shown, drained, readError); + return reply; + } + if (sb == null || encoded < rows.size()) { + // never in JSON mode (fewer than minRows rows, or the columns were dropped late): encode every pulled row now + sb = new StringBuilder(); + meta = new StringBuilder(); + encoded = 0; + for (final Result r : rows) { + final int rowStart = sb.length(), metaStart = meta.length(); + try { + if (encoded > 0) { + sb.append(','); + meta.append(','); + } + NativeCodec.appendRow(sb, meta, r); + } catch (final Throwable t) { + sb.setLength(rowStart); + meta.setLength(metaStart); + readError = ArcadeNative.errorJson(t); + drained = true; + break; + } + encoded++; + } + reply.drained = drained; + } + final long shownJson = reply.drained ? 0 : id; + final StringBuilder out = new StringBuilder("{\"ok\":true,\"rs\":").append(shownJson).append(",\"rows\":[").append(sb) + .append("],\"meta\":[").append(meta).append("],\"drained\":").append(reply.drained); + if (readError != null) + out.append(",\"rerr\":").append(readError); + out.append('}'); + reply.rows = encoded; + reply.json = jsonBytes(out.toString(), sink); + return reply; + } finally { + if (sink != null) + NativeCodec.SINK.remove(); + } + } + + /** The reply text as bytes; with a float32 tail: the text ending ,"bin":N}, a NUL, padding to 4 bytes and the floats (see ArcadeNative.withFloatTail). */ + static byte[] jsonBytes(final String text, final NativeCodec.FloatSink sink) { + if (sink == null || sink.n == 0) + return text.getBytes(StandardCharsets.UTF_8); + final byte[] b = (text.substring(0, text.length() - 1) + ",\"bin\":" + sink.n + "}").getBytes(StandardCharsets.UTF_8); + final int off = (b.length + 1 + 3) & ~3; + final ByteBuffer bb = ByteBuffer.allocate(off + 4 * sink.n); + bb.put(b); + bb.order(ByteOrder.nativeOrder()).position(off).asFloatBuffer().put(sink.a, 0, sink.n); + return bb.array(); + } + + // ---- the columns + + private static int align8(final long n) { + return (int) ((n + 7) & ~7L); + } + + private static boolean yearFits(final int year) { + return year >= 1 && year <= 9999; + } + + private static long micros(final LocalDateTime dt) { + if (!yearFits(dt.getYear())) + throw new Unsupported("year"); + return dt.toEpochSecond(ZoneOffset.UTC) * 1_000_000L + dt.getNano() / 1000; + } + + private static long millisToMicros(final long millis) { + return micros(LocalDateTime.ofEpochSecond(Math.floorDiv(millis, 1000L), Math.floorMod(millis, 1000) * 1_000_000, ZoneOffset.UTC)); + } + + private static long instantMicros(final Instant i) { + return micros(LocalDateTime.ofEpochSecond(i.getEpochSecond(), i.getNano(), ZoneOffset.UTC)); + } + + private static char kindOf(final Object v) { + if (v instanceof String || v instanceof Character) + return 's'; + if (v instanceof Integer || v instanceof Long || v instanceof Short || v instanceof Byte) + return 'q'; + if (v instanceof Double || v instanceof Float) + return 'd'; + if (v instanceof Boolean) + return 'b'; + if (v instanceof float[]) + return 'v'; + if (v instanceof LocalDate) + return 'D'; + if (v instanceof LocalDateTime || v instanceof Date || v instanceof Calendar) + return 'T'; + if (v instanceof Instant || v instanceof ZonedDateTime || v instanceof OffsetDateTime) + return 'Z'; + if (v instanceof Identifiable && !(v instanceof Document)) + return 'R'; + throw new Unsupported(v.getClass().getName()); + } + + /** One column under construction. Cell arrays are zero-filled, so a row without a value costs nothing but its bit and, for strings and vectors, its offset. */ + static final class Col { + final String name; + int idx; + char kind; // 0 until the first value + int rows; // cells so far + long[] l; + double[] d; + byte[] b; + byte[] blob; + int blobLen; + int[] off; + float[] fl; + int flLen; + int dim = -1; + boolean sameDim = true; + byte[] pres = new byte[8], val = new byte[8]; + int nPres, nVal; + // layout, set by Frame + int dataOff, auxOff, valOff, presOff; + char wireKind; + boolean hasVal, hasPres; + + Col(final String name, final int rowsBefore) { + this.name = name; + this.rows = rowsBefore; + } + + private void room() { + final int need = rows + 1; + if ((pres.length << 3) < need) { + final int bytes = Math.max(pres.length * 2, (need + 7) >> 3); + pres = Arrays.copyOf(pres, bytes); + val = Arrays.copyOf(val, bytes); + } + switch (kind) { + case 'q', 'D', 'T', 'Z' -> { + if (l.length < need) + l = Arrays.copyOf(l, Math.max(need, l.length * 2)); + } + case 'd' -> { + if (d.length < need) + d = Arrays.copyOf(d, Math.max(need, d.length * 2)); + } + case 'b' -> { + if (b.length < need) + b = Arrays.copyOf(b, Math.max(need, b.length * 2)); + } + case 's', 'R', 'v' -> { + if (off.length < need + 1) + off = Arrays.copyOf(off, Math.max(need + 1, off.length * 2)); + } + default -> { + } + } + } + + private void init(final char k) { + kind = k; + final int cap = Math.max(16, rows + 16); + switch (k) { + case 'q', 'D', 'T', 'Z' -> l = new long[cap]; + case 'd' -> d = new double[cap]; + case 'b' -> b = new byte[cap]; + case 's', 'R' -> { + off = new int[cap + 1]; + blob = new byte[Math.max(64, cap * 8)]; + } + default -> { // v + off = new int[cap + 1]; + fl = new float[Math.max(64, cap * 8)]; + } + } + } + + /** Adds the cell of row r (== rows) and returns the bytes it brings. */ + long put(final Object v, final int r) { + room(); + pres[r >> 3] |= (byte) (1 << (r & 7)); + nPres++; + if (v == null) { + gap(); + return 0; + } + final char k = kindOf(v); + if (kind == 0) + init(k); + else if (kind != k) + throw new Unsupported("column " + name + " mixes kinds " + kind + " and " + k); + room(); + val[r >> 3] |= (byte) (1 << (r & 7)); + nVal++; + long added = 8; + switch (k) { + case 'q' -> l[rows] = ((Number) v).longValue(); + case 'd' -> d[rows] = ((Number) v).doubleValue(); + case 'b' -> { + b[rows] = (byte) (((Boolean) v) ? 1 : 0); + added = 1; + } + case 'D' -> { + final LocalDate date = (LocalDate) v; + if (!yearFits(date.getYear())) + throw new Unsupported("year"); + l[rows] = date.toEpochDay(); + } + case 'T' -> l[rows] = v instanceof LocalDateTime dt ? micros(dt) : v instanceof Date date ? millisToMicros(date.getTime()) + : millisToMicros(((Calendar) v).getTimeInMillis()); + case 'Z' -> l[rows] = v instanceof Instant i ? instantMicros(i) : v instanceof ZonedDateTime z ? instantMicros(z.toInstant()) + : instantMicros(((OffsetDateTime) v).toInstant()); + case 's', 'R' -> added = 4 + putString(k == 'R' ? ((Identifiable) v).getIdentity().toString() : String.valueOf(v)); + default -> { + final float[] a = (float[]) v; + if (dim < 0) + dim = a.length; + else if (dim != a.length) + sameDim = false; + if (flLen + a.length > fl.length) + fl = Arrays.copyOf(fl, Math.max(fl.length * 2, flLen + a.length)); + System.arraycopy(a, 0, fl, flLen, a.length); + flLen += a.length; + off[rows + 1] = flLen; + added = 4 + 4L * a.length; + } + } + rows++; + return added; + } + + /** A string cell of a column without kind logic or bitmaps (the record meta). */ + int putPlain(final String s) { + if (off.length < rows + 2) + off = Arrays.copyOf(off, Math.max(rows + 2, off.length * 2)); + final int added = putString(s); + rows++; + return added + 4; + } + + /** A row with the property null or absent: a zero cell. */ + void gap() { + if (kind == 's' || kind == 'R' || kind == 'v') + off[rows + 1] = kind == 'v' ? flLen : blobLen; + rows++; + } + + void absent() { + room(); + gap(); + } + + private int putString(final String s) { + final int len = s.length(); + if (blobLen + 3L * len > blob.length) + blob = Arrays.copyOf(blob, (int) Math.min(Integer.MAX_VALUE - 8L, Math.max(blob.length * 2L, blobLen + 3L * len))); + int p = blobLen; + final byte[] o = blob; + for (int i = 0; i < len; i++) { + final char c = s.charAt(i); + if (c < 0x80) + o[p++] = (byte) c; + else if (c < 0x800) { + o[p++] = (byte) (0xC0 | (c >> 6)); + o[p++] = (byte) (0x80 | (c & 0x3F)); + } else if (c >= 0xD800 && c <= 0xDFFF) { + if (c <= 0xDBFF && i + 1 < len && Character.isLowSurrogate(s.charAt(i + 1))) { + final int cp = Character.toCodePoint(c, s.charAt(++i)); + o[p++] = (byte) (0xF0 | (cp >> 18)); + o[p++] = (byte) (0x80 | ((cp >> 12) & 0x3F)); + o[p++] = (byte) (0x80 | ((cp >> 6) & 0x3F)); + o[p++] = (byte) (0x80 | (cp & 0x3F)); + } else + throw new Unsupported("unpaired surrogate"); + } else { + o[p++] = (byte) (0xE0 | (c >> 12)); + o[p++] = (byte) (0x80 | ((c >> 6) & 0x3F)); + o[p++] = (byte) (0x80 | (c & 0x3F)); + } + } + final int added = p - blobLen; + blobLen = p; + off[rows + 1] = p; + return added; + } + } + + /** The rows of one batch as columns, filled row by row. */ + static final class Builder { + final ArrayList cols = new ArrayList<>(); + final Map byName = new HashMap<>(); + Col metaRid, metaType; + byte[] metaKind = new byte[64]; + int n; + long bytes; + long cells; + final int[] seq = new int[MAX_COLS + 1]; + final Map orders = new java.util.LinkedHashMap<>(); // rows whose property order is not the column order: the order, as column indexes + + void add(final Result row) { + final int r = n; + int last = -1, count = 0; + boolean inOrder = true; + for (final String name : row.getPropertyNames()) { + Col c = byName.get(name); + if (c == null) { + if (cols.size() >= MAX_COLS) + throw new Unsupported("columns"); + c = new Col(name, r); + c.idx = cols.size(); + cols.add(c); + byName.put(name, c); + } + if (c.idx <= last) + inOrder = false; + last = c.idx; + seq[count++] = c.idx; + bytes += c.put(row.getProperty(name), r); + cells++; + } + for (final Col c : cols) + if (c.rows == r) + c.absent(); + if (!inOrder) + orders.put(r, Arrays.copyOf(seq, count)); + byte k = 0; + String rid = "", type = ""; + if (row.isElement()) { + final Optional element = row.getElement(); + if (element.isPresent()) { + final Document d = element.get(); + k = (byte) (d instanceof Vertex ? 'v' : d instanceof Edge ? 'e' : 'd'); + rid = d.getIdentity() == null ? "" : d.getIdentity().toString(); + type = d.getTypeName(); + } + } + if (k != 0 && metaRid == null) { + metaRid = new Col("@rid", r); + metaType = new Col("@type", r); + metaRid.kind = metaType.kind = 's'; + metaRid.off = new int[Math.max(16, r + 16) + 1]; + metaType.off = new int[Math.max(16, r + 16) + 1]; + metaRid.blob = new byte[256]; + metaType.blob = new byte[256]; + metaRid.pres = metaType.pres = null; // the meta columns carry no bitmaps (Frame.layout reads this) + } + if (metaKind.length <= r) + metaKind = Arrays.copyOf(metaKind, metaKind.length * 2); + metaKind[r] = k; + if (metaRid != null) { + bytes += 1 + metaRid.putPlain(rid) + metaType.putPlain(type); + } + n++; + } + + /** A table that is mostly holes is bigger and slower as columns than as rows. */ + boolean sparse() { + return (cols.size() > 1 && cells * 4 < (long) n * cols.size()) || orders.size() > n / 2; + } + + Frame finish(final long rs, final boolean drained, final String rerr) { + return new Frame(this, rs, drained, rerr); + } + } + + /** The laid-out frame: sizes and offsets decided, ready to be written into a buffer of exactly total bytes. */ + static final class Frame { + final Builder b; + final long rs; + final boolean drained; + final byte[] hdr; + final int hdrLen; + final int total; + int pos = 8; // offsets within the data section; 0 means none + int metaKindAt; + + Frame(final Builder b, final long rs, final boolean drained, final String rerr) { + this.b = b; + this.rs = rs; + this.drained = drained; + final int n = b.n; + for (final Col c : b.cols) + layout(c, n); + int metaKind = 0; + if (b.metaRid != null) { + metaKind = pos; + metaKindAt = pos; + pos += align8(n); + layout(b.metaRid, n); + layout(b.metaType, n); + } + final StringBuilder sb = new StringBuilder("{\"c\":["); + boolean first = true; + for (final Col c : b.cols) { + if (!first) + sb.append(','); + first = false; + sb.append('['); + NativeCodec.quote(sb, c.name); + sb.append(",\"").append(c.wireKind).append("\",").append(c.wireKind == 'f' ? c.dim : 0).append(',').append(c.dataOff).append(',') + .append(c.auxOff).append(',').append(c.valOff).append(',').append(c.presOff).append(']'); + } + sb.append("],\"m\":"); + if (b.metaRid == null) + sb.append('0'); + else + sb.append('[').append(metaKind).append(',').append(b.metaRid.dataOff).append(',').append(b.metaRid.auxOff).append(',') + .append(b.metaType.dataOff).append(',').append(b.metaType.auxOff).append(']'); + if (!b.orders.isEmpty()) { + sb.append(",\"o\":["); + boolean firstOrder = true; + for (final Map.Entry e : b.orders.entrySet()) { + if (!firstOrder) + sb.append(','); + firstOrder = false; + sb.append('[').append(e.getKey()).append(",["); + for (int i = 0; i < e.getValue().length; i++) + sb.append(i > 0 ? "," : "").append(e.getValue()[i]); + sb.append("]]"); + } + sb.append(']'); + } + if (rerr != null) + sb.append(",\"rerr\":").append(rerr); + sb.append('}'); + while ((sb.length() & 7) != 0) + sb.append(' '); + hdr = sb.toString().getBytes(StandardCharsets.UTF_8); + hdrLen = hdr.length; + total = HEADER + hdrLen + pos; + } + + private void layout(final Col c, final int n) { + c.hasVal = c.nVal < n; + c.hasPres = c.nPres < n; + if (c.pres == null) { // meta columns carry no bitmaps + c.hasVal = c.hasPres = false; + } + if (c.hasVal) { + c.valOff = pos; + pos += align8((n + 7) >> 3); + } + if (c.hasPres) { + c.presOff = pos; + pos += align8((n + 7) >> 3); + } + char k = c.kind == 0 ? 'n' : c.kind; + if (k == 'v' && !c.hasVal && c.sameDim && c.dim >= 1) + k = 'f'; + c.wireKind = k; + switch (k) { + case 'q', 'd', 'D', 'T', 'Z' -> { + c.dataOff = pos; + pos += align8(8L * n); + } + case 'b' -> { + c.dataOff = pos; + pos += align8(n); + } + case 's', 'R' -> { + c.dataOff = pos; + pos += align8(4L * (n + 1)); + c.auxOff = pos; + pos += align8(c.blobLen); + } + case 'v' -> { + c.dataOff = pos; + pos += align8(4L * (n + 1)); + c.auxOff = pos; + pos += align8(4L * c.flLen); + } + case 'f' -> { + c.dataOff = pos; + pos += align8(4L * n * c.dim); + } + default -> { + } + } + } + + void write(final ByteBuffer out) { + final ByteBuffer bb = out; + bb.order(ByteOrder.nativeOrder()); + final int base = HEADER + hdrLen; + bb.putInt(0, MAGIC); + bb.putInt(4, hdrLen); + bb.putLong(8, total); + bb.putLong(16, rs); + bb.putInt(24, b.n); + bb.putInt(28, drained ? 1 : 0); + bb.put(HEADER, hdr, 0, hdrLen); + bb.put(base, ZEROS, 0, 8); + final int n = b.n; + for (final Col c : b.cols) + writeCol(bb, base, c, n); + if (b.metaRid != null) { + final int at = base + metaKindAt; + bb.put(at, b.metaKind, 0, n); + pad(bb, at + n); + writeCol(bb, base, b.metaRid, n); + writeCol(bb, base, b.metaType, n); + } + } + + private static void pad(final ByteBuffer bb, final int end) { + final int p = align8(end) - end; + if (p > 0) + bb.put(end, ZEROS, 0, p); + } + + private static void writeCol(final ByteBuffer bb, final int base, final Col c, final int n) { + final int bitmap = (n + 7) >> 3; + if (c.hasVal) { + bb.put(base + c.valOff, c.val, 0, bitmap); + pad(bb, base + c.valOff + bitmap); + } + if (c.hasPres) { + bb.put(base + c.presOff, c.pres, 0, bitmap); + pad(bb, base + c.presOff + bitmap); + } + final int at = base + c.dataOff; + switch (c.wireKind) { + case 'q', 'd', 'D', 'T', 'Z' -> { + final ByteBuffer s = bb.slice(at, 8 * n).order(ByteOrder.nativeOrder()); + if (c.wireKind == 'd') + s.asDoubleBuffer().put(c.d, 0, n); + else + s.asLongBuffer().put(c.l, 0, n); + } + case 'b' -> { + bb.put(at, c.b, 0, n); + pad(bb, at + n); + } + case 's', 'R' -> { + bb.slice(at, 4 * (n + 1)).order(ByteOrder.nativeOrder()).asIntBuffer().put(c.off, 0, n + 1); + pad(bb, at + 4 * (n + 1)); + bb.put(base + c.auxOff, c.blob, 0, c.blobLen); + pad(bb, base + c.auxOff + c.blobLen); + } + case 'v' -> { + bb.slice(at, 4 * (n + 1)).order(ByteOrder.nativeOrder()).asIntBuffer().put(c.off, 0, n + 1); + pad(bb, at + 4 * (n + 1)); + bb.slice(base + c.auxOff, 4 * c.flLen).order(ByteOrder.nativeOrder()).asFloatBuffer().put(c.fl, 0, c.flLen); + pad(bb, base + c.auxOff + 4 * c.flLen); + } + case 'f' -> { + bb.slice(at, 4 * n * c.dim).order(ByteOrder.nativeOrder()).asFloatBuffer().put(c.fl, 0, n * c.dim); + pad(bb, at + 4 * n * c.dim); + } + default -> { + } + } + } + } +} diff --git a/native-spike/src/main/java/org/humemai/nativelib/NativeCodec.java b/native-spike/src/main/java/org/humemai/nativelib/NativeCodec.java index 189dab7cf2f..498301b2877 100644 --- a/native-spike/src/main/java/org/humemai/nativelib/NativeCodec.java +++ b/native-spike/src/main/java/org/humemai/nativelib/NativeCodec.java @@ -92,7 +92,7 @@ static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int return drained; } - private static void appendRow(final StringBuilder sb, final StringBuilder meta, final Result row) { + static void appendRow(final StringBuilder sb, final StringBuilder meta, final Result row) { sb.append('{'); boolean first = true; for (final String name : row.getPropertyNames()) {