From 76f94b2bb8a66b574f0222fa66efaad082a76cd5 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Mon, 9 Nov 2026 17:44:19 +0900 Subject: [PATCH 1/2] JVM to_json_list: a FLOAT and a float32 float[] arrive as the exact double, as on the native wheel (refs #314) (#320) RowBatcher.keepNonFinite now widens a Float to a Double and a float[] to a double[] before the engine's JSONObject writes it, so the JSON text is the exact float64 widening of the float32 (0.10000000149011612, not the shortest float32 text 0.1). NaN and infinity keep their tag. tests/test_cross_wheel_golden.py is one fixed table with committed golden literals for to_list, to_json_list and to_columns (FLOAT, DOUBLE, float32 float[] of 3 and 9, DECIMAL, LONG beyond 2**53, DATE, DATETIME, NULL, embedded list), run unchanged on both wheels. Co-authored-by: Claude Sonnet 5.5 (cherry picked from commit 551dee97c7a46c3a51f1f84fe40bb2824704dedb) Co-Authored-By: Claude Sonnet 5.5 --- src/java/com/arcadedb/python/RowBatcher.java | 37 ++-- tests/test_cross_wheel_golden.py | 215 +++++++++++++++++++ 2 files changed, 240 insertions(+), 12 deletions(-) create mode 100644 tests/test_cross_wheel_golden.py diff --git a/src/java/com/arcadedb/python/RowBatcher.java b/src/java/com/arcadedb/python/RowBatcher.java index b1548348..150299ec 100644 --- a/src/java/com/arcadedb/python/RowBatcher.java +++ b/src/java/com/arcadedb/python/RowBatcher.java @@ -53,23 +53,36 @@ private RowBatcher() { * JSON has no NaN or infinity, and the engine's writers turn them into null. This returns {@code value} with every * NaN, Infinity and -Infinity (a Double or Float, or an element of a float[], a double[], a Collection or a Map, at any * depth) replaced by the tag {@code {"\u0001": ["f", "NaN"]}} (the text is {@code Double.toString}), which Python turns - * back into a float. A value with nothing non-finite is returned as it is, with no copy, so the common row costs one - * type test per value. + * back into a float. A Float, and each element of a float[], is also widened to a double (a float[] becomes a double[]), + * so the JSON text is the exact double and Python reads the same float that to_list() returns (#314). A value that is + * neither a Float nor a float[] and has nothing non-finite is returned as it is, with no copy, so the common row costs + * one type test per value. */ static Object keepNonFinite(final Object value) { - if (value instanceof Double || value instanceof Float) { - final double d = ((Number) value).doubleValue(); + if (value instanceof Double) { + final double d = (Double) value; return Double.isFinite(d) ? value : tag(d); } + if (value instanceof Float) { + // widened to a Double: the JSON writer prints a Float as the shortest float32 text ("0.1"), which Python reads as a + // different double than the exact widening of the float32 that to_list() and the native build return + final double d = ((Float) value).doubleValue(); + return Double.isFinite(d) ? (Object) d : tag(d); + } if (value instanceof float[] floats) { - for (final float f : floats) - if (!Float.isFinite(f)) { - final List out = new ArrayList<>(floats.length); - for (final float x : floats) - out.add(Float.isFinite(x) ? (Object) x : tag(x)); - return out; - } - return value; + final double[] wide = new double[floats.length]; + boolean finite = true; + for (int i = 0; i < floats.length; i++) { + wide[i] = floats[i]; + if (!Double.isFinite(wide[i])) + finite = false; + } + if (finite) + return wide; + final List out = new ArrayList<>(wide.length); + for (final double x : wide) + out.add(Double.isFinite(x) ? (Object) x : tag(x)); + return out; } if (value instanceof double[] doubles) { for (final double d : doubles) diff --git a/tests/test_cross_wheel_golden.py b/tests/test_cross_wheel_golden.py new file mode 100644 index 00000000..a5e62209 --- /dev/null +++ b/tests/test_cross_wheel_golden.py @@ -0,0 +1,215 @@ +"""One fixed table, committed golden literals, run on both wheels: the two builds return the same values (#314). + +``tests/`` runs against each installed wheel, but until now a value was checked against one literal in a test +that only one wheel could run, so a difference between the wheels stayed invisible. This file has no branch on +``arcadedb.__backend__``: the golden literals below are asserted on ``arcadedb`` (native) and on ``arcadedb-jvm`` +alike, so a literal that holds on both IS the parity check. + +The table holds a FLOAT, a DOUBLE, two float32 ``float[]`` (3 values, which the native build writes element by +element, and 9, which it sends as a float32 tail), a DECIMAL, a LONG beyond 2**53, a DATE, a DATETIME, a NULL and an +embedded list. Every float32 input is written as the exact float64 widening of that float32, so the stored value is +the literal. ``to_json_list()`` of a FLOAT or a float32 vector used to differ: the JVM build printed the shortest +float32 text (``0.1``), which Python reads as a different double than the exact widening +(``0.10000000149011612``) that the native build and ``to_list()`` return. +""" + +from datetime import date, datetime +from decimal import Decimal + +import numpy as np +import pytest + +SELECT = "SELECT n, f, d, v3, v9, dec, big, day, ts, maybe, lst FROM Golden ORDER BY n" + +# float32 values written as the float64 they widen to, so each is exact in a FLOAT and in a float32 vector +F32_TENTH = 0.10000000149011612 # float32(0.1) +F32_THIRD = 0.3333333432674408 # float32(1/3) +F32_MILLI = 0.0010000000474974513 # float32(0.001) +F32_SEVEN_TENTHS = 0.699999988079071 # float32(0.7) +F32_NINE_TENTHS = 0.8999999761581421 # float32(0.9) + +V9_ROW0 = [ + 0.10000000149011612, + 0.20000000298023224, + 0.30000001192092896, + 0.4000000059604645, + 0.5, + 0.6000000238418579, + 0.699999988079071, + 0.800000011920929, + 0.8999999761581421, +] +V9_ROW1 = [ + F32_MILLI, + F32_SEVEN_TENTHS, + F32_NINE_TENTHS, + 0.5, + F32_THIRD, + F32_TENTH, + 1.0, + 2.0, + -2.75, +] + +ROWS = [ + { + "n": 0, + "f": F32_TENTH, + "d": 0.1, + "v3": [F32_TENTH, F32_THIRD, 0.5], + "v9": V9_ROW0, + "dec": Decimal("1.50"), + "big": 2**53 + 1, + "day": date(2026, 1, 2), + "ts": datetime(2026, 1, 2, 3, 4, 5, 678000), + "maybe": 0.25, + "lst": [1, "a", 2.5], + }, + { + "n": 1, + "f": 0.5, + "d": 1 / 3, + "v3": [0.5, 1.0, -2.75], + "v9": V9_ROW1, + "dec": Decimal("-0.25"), + "big": -(2**53) - 1, + "day": date(1969, 12, 31), + "ts": datetime(1969, 12, 31, 23, 59, 59, 999000), + "maybe": None, + "lst": ["x"], + }, +] + +# to_list: the Python types of the values as stored +GOLDEN_LIST = ROWS + +# to_json_list: JSON-native values; DATE and DATETIME as epoch milliseconds, DECIMAL as a float +GOLDEN_JSON = [ + { + "n": 0, + "f": F32_TENTH, + "d": 0.1, + "v3": [F32_TENTH, F32_THIRD, 0.5], + "v9": V9_ROW0, + "dec": 1.5, + "big": 9007199254740993, + "day": 1767312000000, + "ts": 1767323045678, + "maybe": 0.25, + "lst": [1, "a", 2.5], + }, + { + "n": 1, + "f": 0.5, + "d": 1 / 3, + "v3": [0.5, 1.0, -2.75], + "v9": V9_ROW1, + "dec": -0.25, + "big": -9007199254740993, + "day": -86400000, + "ts": -1, + "maybe": None, + "lst": ["x"], + }, +] + +# to_columns: one array per column, the numpy dtype is part of the golden +GOLDEN_COLUMNS = { + "n": np.array([0, 1], dtype=np.int64), + "f": np.array([F32_TENTH, 0.5], dtype=np.float64), + "d": np.array([0.1, 1 / 3], dtype=np.float64), + "v3": np.array([[F32_TENTH, F32_THIRD, 0.5], [0.5, 1.0, -2.75]], dtype=np.float32), + "v9": np.array([V9_ROW0, V9_ROW1], dtype=np.float32), + "big": np.array([9007199254740993, -9007199254740993], dtype=np.int64), + "day": np.array(["2026-01-02", "1969-12-31"], dtype="datetime64[ms]"), + "ts": np.array( + ["2026-01-02T03:04:05.678", "1969-12-31T23:59:59.999"], dtype="datetime64[ms]" + ), + "maybe": np.array([0.25, np.nan], dtype=np.float64), +} + + +@pytest.fixture +def golden(temp_db): + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE Golden") + for name, type_name in [ + ("n", "INTEGER"), + ("f", "FLOAT"), + ("d", "DOUBLE"), + ("v3", "ARRAY_OF_FLOATS"), + ("v9", "ARRAY_OF_FLOATS"), + ("dec", "DECIMAL"), + ("big", "LONG"), + ("day", "DATE"), + ("ts", "DATETIME"), + ("maybe", "DOUBLE"), + ("lst", "LIST"), + ]: + db.command("sql", f"CREATE PROPERTY Golden.{name} {type_name}") + db.insert_many("Golden", [dict(row) for row in ROWS]) + return db + + +def _json_rows(db): + return db.query("sql", SELECT).to_json_list() + + +def test_to_json_list_floats_are_the_exact_widening_of_the_float32(golden): + """The float cases of #314: FLOAT, float32 float[] of 3 and of 9 values, next to a DOUBLE and the shortest-text trap.""" + rows = _json_rows(golden) + for have, want in zip(rows, GOLDEN_JSON): + for name in ("f", "d", "v3", "v9", "maybe"): + assert have[name] == want[name], (have["n"], name) + # the shortest text of float32(0.1) is "0.1", which Python reads as a different double + assert rows[0]["f"] != 0.1 + assert rows[0]["f"] == 0.10000000149011612 + + +def test_to_json_list_floats_equal_to_list_floats(golden): + """The two readers of the same rows agree: row["v9"] == to_list()[i]["v9"] (false on the JVM build before #314).""" + as_json = _json_rows(golden) + as_list = golden.query("sql", SELECT).to_list() + for from_json, from_list in zip(as_json, as_list): + for name in ("f", "d", "v3", "v9", "maybe"): + assert from_json[name] == from_list[name], (from_json["n"], name) + + +def test_to_json_list_golden(golden): + assert _json_rows(golden) == GOLDEN_JSON + rows = _json_rows(golden) + assert type(rows[0]["dec"]) is float + assert type(rows[0]["big"]) is int and type(rows[0]["day"]) is int + assert type(rows[0]["ts"]) is int + assert all(type(x) is float for row in rows for x in row["v3"] + row["v9"]) + + +def test_to_list_golden(golden): + rows = golden.query("sql", SELECT).to_list() + assert rows == GOLDEN_LIST + assert type(rows[0]["dec"]) is Decimal + assert type(rows[0]["big"]) is int + assert type(rows[0]["day"]) is date + assert type(rows[0]["ts"]) is datetime + assert all(type(x) is float for row in rows for x in row["v3"] + row["v9"]) + + +def _as_list(column): + return [list(x) if isinstance(x, (list, np.ndarray)) else x for x in column] + + +@pytest.mark.parametrize("name", sorted(GOLDEN_COLUMNS)) +def test_to_columns_golden(golden, name): + column = golden.query("sql", SELECT).to_columns()[name] + want = GOLDEN_COLUMNS[name] + assert isinstance(column, np.ndarray), type(column) + assert column.dtype == want.dtype, (column.dtype, want.dtype) + assert column.shape == want.shape + np.testing.assert_array_equal(column, want) + + +def test_to_columns_decimal_and_list_columns(golden): + columns = golden.query("sql", SELECT).to_columns() + assert list(columns["dec"]) == [Decimal("1.50"), Decimal("-0.25")] + assert all(type(v) is Decimal for v in columns["dec"]) + assert _as_list(columns["lst"]) == [[1, "a", 2.5], ["x"]] From 946dfda18fa713bb92d1171d6a5084cd98f5d4a8 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Mon, 9 Nov 2026 18:05:10 +0900 Subject: [PATCH 2/2] Native columnar readers: infer Arrow types first, skip per-key and per-row steps, one _kind per column (refs #312) (#321) * Native columnar readers: infer Arrow types first, skip per-key and per-row steps, one _kind per column (refs #312) Bindings-only, Python in arcadedb._native, no image rebuild. Output is meant to be identical to before: - columns.to_arrow tries pa.array(values) first and keeps the old per-value rules for a column that inference cannot type (ArrowInvalid, ArrowTypeError) or types differently (temporal types, unsigned ints). - columns._names tests each row's keys against the known ones in C instead of a Python step per key. - columns.to_columns calls _kind once per column instead of twice. - NativeResultSet._fill can hand a reply's rows straight to _drain_dicts, which no longer builds a tuple and a generator step per row. The float32-tail view in codec.loads_tagged is not done: a vector must stay a list of floats for the row API. tests/test_native_columns_reference.py keeps the old rules as a private reference and compares tables, columns, names and the drain on mixed, null-heavy, ragged and vector inputs. No speed is claimed here. Co-Authored-By: Claude Sonnet 5.5 * Native columns: a failed inference of any kind falls to the old rules; tests compare NaN tables by value (refs #312) pa.array([2**63]) raises OverflowError, not ArrowInvalid, so an int above int64 next to floats would have raised where the old explicit-type call fell back to exact text. _inferred now leaves every failed inference to the rules, so the outcome of such a column is decided by the old code. Four reference cases cover it. Table.equals is False for a NaN next to the same NaN, so the comparison is the schema plus the repr of the values. Co-Authored-By: Claude Sonnet 5.5 --------- Co-authored-by: Claude Sonnet 5.5 (cherry picked from commit 6c6a5a12cbca3ad2ae6ab7993023a8c4de39d2e7) Co-Authored-By: Claude Sonnet 5.5 --- src/arcadedb/_native/columns.py | 40 ++- src/arcadedb/_native/database.py | 32 +- tests/conftest.py | 2 + tests/test_native_columns_reference.py | 428 +++++++++++++++++++++++++ 4 files changed, 489 insertions(+), 13 deletions(-) create mode 100644 tests/test_native_columns_reference.py diff --git a/src/arcadedb/_native/columns.py b/src/arcadedb/_native/columns.py index 90f59df6..1d789922 100644 --- a/src/arcadedb/_native/columns.py +++ b/src/arcadedb/_native/columns.py @@ -5,12 +5,16 @@ def _names(rows, columns): + """The column names: ``columns``, or the keys of the rows in the order they are first seen.""" if columns is not None: return list(columns) seen = {} + known = seen.keys() for row in rows: - for k in row: - seen.setdefault(k, None) + # a row whose keys are all known (nearly every row) costs one C-level subset test, not one Python step per key + if not row.keys() <= known: + for k in row: + seen.setdefault(k, None) return list(seen) @@ -84,13 +88,13 @@ def to_columns(rows, columns=None): out = {} for name in _names(rows, columns): values = [row.get(name) for row in rows] - if _kind(values) == "obj": + kind = _kind(values) + if kind == "obj": vectors = _float_vectors(values) if vectors is not None: out[name] = vectors continue has_null = any(v is None for v in values) - kind = _kind(values) if kind == "i": out[name] = ( np.array([np.nan if v is None else v for v in values], dtype=np.float64) @@ -124,12 +128,40 @@ def to_columns(rows, columns=None): return out +def _inferred(values): + """``pa.array(values)`` when pyarrow's own inference gives the array the rules of :func:`to_arrow` would give, else None. + + Inference replaces the per-value Python pass of :func:`_kind` for a column of ints, floats, bools, strings, lists and + other values it types on its own. Two inferred types differ from the rules: a temporal type (the rules write + ``timestamp[ms]`` for a datetime and a date, inference gives ``timestamp[us]`` or ``date32``) and an unsigned integer + (an int above int64 is an error under the rules). A column that inference cannot type is also left to the rules, + whatever the error (``ArrowInvalid``, ``ArrowTypeError``, and ``OverflowError`` for an int above int64, which the + explicit-type call turns into ``ArrowInvalid``), so a failed inference never decides the outcome: the rules have the + fallbacks (exact text for an int beyond 2**53 next to floats, for a Decimal wider than decimal256, for a value of mixed + types). + """ + import pyarrow as pa + + try: + array = pa.array(values) + except Exception: # the rules decide the outcome of a column inference cannot type + return None + kind = array.type + if pa.types.is_temporal(kind) or pa.types.is_unsigned_integer(kind): + return None + return array + + def to_arrow(rows, columns=None): import pyarrow as pa arrays = {} for name in _names(rows, columns): values = [row.get(name) for row in rows] + inferred = _inferred(values) + if inferred is not None: + arrays[name] = inferred + continue kind = _kind(values) if kind == "i": arrays[name] = pa.array(values, type=pa.int64()) diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index f30dc526..924d5a17 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -248,8 +248,13 @@ def _readable(self) -> bool: "with block), so the rows it had not returned are gone. Run the query again to read them." ) - def _fill(self) -> bool: - """Fetch the next batch into the buffer; False once the cursor is drained.""" + def _fill(self, into: Optional[list] = None) -> bool: + """Fetch the next batch into the buffer; False once the cursor is drained. + + ``into``: a list that takes the rows (without their meta) of the batch when the buffer is empty, so a drain that + wants only the rows does not build one tuple per row to take them apart again. The buffer then stays empty, and the + answer is the same: whether the batch had a row (or an error is pending). + """ if self._pending_error is not None: err, self._pending_error = self._pending_error, None self._more = False @@ -279,10 +284,16 @@ def _fill(self) -> bool: raise ArcadeDBError( f"Reading the result set failed: {reply.get('cls')}: {reply.get('error')}" ) - self._buf.extend(zip(reply["rows"], reply["meta"])) + if into is not None and not self._buf: + batch = reply["rows"] + into.extend(batch) + has_rows = bool(batch) + else: + self._buf.extend(zip(reply["rows"], reply["meta"])) + has_rows = bool(self._buf) self._more = not reply["drained"] self._pending_error = reply.get("rerr") - return bool(self._buf) or self._pending_error is not None + return has_rows or self._pending_error is not None def _finish(self) -> None: self._exhausted = True @@ -334,11 +345,14 @@ def _drain_dicts(self) -> List[Dict[str, Any]]: out: List[Dict[str, Any]] = [] if not self._readable(): return out - while True: - out.extend(p for p, _ in self._buf) - self._buf.clear() - if not self._fill(): - break + buf = self._buf + if ( + buf + ): # rows read ahead (the first reply, or rows a caller has not taken): the only rows that carry a tuple + out.extend([p for p, _ in buf]) + buf.clear() + while self._fill(out): + pass self._finish() return out diff --git a/tests/conftest.py b/tests/conftest.py index ec3147c7..eaa3c695 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -83,6 +83,8 @@ collect_ignore.append("test_native_point_op_glue.py") # the address table of append_samples (arcadedb_n_async_bin), observed on the native library collect_ignore.append("test_native_timeseries_binary.py") + # arcadedb._native.columns builds the columnar readers' tables from the rows; the file compares it with the rules it replaced (#312) + collect_ignore.append("test_native_columns_reference.py") if not os.path.isfile( os.path.join(os.path.dirname(os.path.dirname(__file__)), "docs", "index.md") ): diff --git a/tests/test_native_columns_reference.py b/tests/test_native_columns_reference.py new file mode 100644 index 00000000..f1f68bad --- /dev/null +++ b/tests/test_native_columns_reference.py @@ -0,0 +1,428 @@ +"""The native columnar readers give the same table as the rules they replaced (refs #312). + +``arcadedb._native.columns`` infers an Arrow column with ``pa.array(values)`` first and falls back to the per-value +rules (``_kind``) when inference cannot type the column or types it differently; ``to_columns`` calls ``_kind`` once +per column; ``_names`` skips the per-key step for a row whose keys are known; ``NativeResultSet._drain_dicts`` takes +the rows of a reply without one tuple per row. None of that may change a value, a dtype or a column order. + +The reference functions below are the code as it was before: the rules of ``_kind``, ``_names``, ``to_columns`` and +``to_arrow``, copied. ``_ms`` and ``_float_vectors`` were not changed and are shared with the module. Each case builds +rows, runs both and compares the outcome, an exception included (a column that raised before raises the same class now). +Pure Python on the rows: no database is needed except for the drain test at the end. + +conftest.py collects this file on the native backend only (the module lives in the native package). +""" + +import random +from datetime import date, datetime, timedelta, timezone +from decimal import Decimal + +import numpy as np +import pyarrow as pa +import pytest + +from arcadedb._native import columns as new +from arcadedb._native.columns import _float_vectors, _ms + +# --- the reference: the rules before this change -------------------------------------------------------------------------- + + +def _names_ref(rows, columns): + if columns is not None: + return list(columns) + seen = {} + for row in rows: + for k in row: + seen.setdefault(k, None) + return list(seen) + + +def _kind_ref(values): + kinds = set() + for v in values: + if v is None: + continue + t = type(v) + if t is bool: + kinds.add("b") + elif t is int: + kinds.add("i") + elif t is float: + kinds.add("f") + elif t is datetime or t is date: + kinds.add("dt") + elif t is Decimal: + kinds.add("dec") + else: + kinds.add("obj") + if not kinds: + return "obj" + if len(kinds) == 1: + return kinds.pop() + if kinds == {"i", "f"}: + return "f" + return "obj" + + +def _to_columns_ref(rows, columns=None): + out = {} + for name in _names_ref(rows, columns): + values = [row.get(name) for row in rows] + if _kind_ref(values) == "obj": + vectors = _float_vectors(values) + if vectors is not None: + out[name] = vectors + continue + has_null = any(v is None for v in values) + kind = _kind_ref(values) + if kind == "i": + out[name] = ( + np.array([np.nan if v is None else v for v in values], dtype=np.float64) + if has_null + else np.array(values, dtype=np.int64) + ) + elif kind == "f": + out[name] = np.array( + [np.nan if v is None else v for v in values], dtype=np.float64 + ) + elif kind == "b": + out[name] = values if has_null else np.array(values, dtype=bool) + elif kind == "dt": + out[name] = np.array( + [ + ( + np.datetime64("NaT", "ms") + if v is None + else np.datetime64(_ms(v), "ms") + ) + for v in values + ], + dtype="datetime64[ms]", + ) + elif kind == "dec": + arr = np.empty(len(values), dtype=object) + arr[:] = values + out[name] = arr + else: + out[name] = values + return out + + +def _to_arrow_ref(rows, columns=None): + arrays = {} + for name in _names_ref(rows, columns): + values = [row.get(name) for row in rows] + kind = _kind_ref(values) + if kind == "i": + arrays[name] = pa.array(values, type=pa.int64()) + elif kind == "f": + try: + arrays[name] = pa.array(values, type=pa.float64()) + except pa.ArrowInvalid: + arrays[name] = pa.array( + [None if v is None else str(v) for v in values], type=pa.string() + ) + elif kind == "b": + arrays[name] = pa.array(values, type=pa.bool_()) + elif kind == "dt": + arrays[name] = pa.array( + [None if v is None else _ms(v) for v in values], type=pa.timestamp("ms") + ) + elif kind == "dec": + try: + arrays[name] = pa.array(values) + except pa.ArrowInvalid: + arrays[name] = pa.array( + [None if v is None else str(v) for v in values], type=pa.string() + ) + else: + try: + arrays[name] = pa.array([None if v is None else v for v in values]) + except (pa.ArrowInvalid, pa.ArrowTypeError): + arrays[name] = pa.array( + [None if v is None else str(v) for v in values], type=pa.string() + ) + return pa.table(arrays) + + +# --- the inputs ----------------------------------------------------------------------------------------------------------- + +BIG = 2**53 + 1 +STAMP = datetime(2026, 1, 2, 3, 4, 5, 678901) +AWARE = datetime(2026, 1, 2, 12, 0, 0, tzinfo=timezone(timedelta(hours=9))) +WIDE40 = Decimal("1234567890123456789012345678901234567890.5") # decimal256 +WIDE80 = Decimal("1" * 80) # wider than decimal256: exact text + +# one column per case, each a list of values that go in rows as {"c": value} +COLUMNS = { + "ints": [1, 2, 3], + "ints-null": [1, None, 3], + "ints-int64-edges": [-(2**63), 2**63 - 1, 0], + "int-above-int64": [2**63, 1], + "int-above-uint64": [2**64, 1], + "int-below-int64": [-(2**63) - 1, 1], + "floats": [0.5, 1.5, float("nan")], + "floats-null": [0.5, None, float("inf")], + "ints-and-floats": [1, 2.5, None, 3], + "int-beyond-2-53-and-floats": [BIG, 0.5], + "big-int-beyond-2-53-and-floats": [2**60, 1.5, None], + "int-above-int64-and-floats": [2**63, 0.5], + "int-above-uint64-and-floats": [2**64, 1.5, None], + "int-below-int64-and-floats": [-(2**63) - 1, 0.5], + "int-above-int64-and-string": [2**63, "x"], + "bools": [True, False, True], + "bools-null": [True, None, False], + "bool-and-int": [True, 1, 0], + "bool-and-float": [True, 0.5], + "strings": ["a", "bb", ""], + "strings-null": ["a", None, "c"], + "string-and-int": ["seven", 1], + "all-null": [None, None, None], + "single-null": [None], + "datetimes": [STAMP, datetime(1969, 12, 31, 23, 59, 59, 999000)], + "datetimes-null": [STAMP, None], + "dates": [date(2026, 1, 2), date(1969, 12, 31)], + "date-and-datetime": [date(2026, 1, 2), STAMP, None], + "aware-datetime": [AWARE, STAMP], + "datetime-and-string": [STAMP, "x"], + "decimals": [Decimal("1.50"), Decimal("-0.25")], + "decimals-null": [Decimal("2"), None, Decimal("30")], + "decimal-40-digits": [WIDE40, Decimal("1")], + "decimal-80-digits": [WIDE80, Decimal("1")], + "decimal-and-int": [Decimal("1.5"), 2], + "vectors": [[0.5, 1.5], [0.25, 2.0]], + "vectors-null": [[0.5, 1.5], None, [0.25, 2.0]], + "vector-not-float32": [[0.1, 0.2], [0.3, 0.4]], + "vector-ragged": [[0.5], [0.5, 1.5]], + "lists-of-ints": [[1, 2], [3]], + "lists-mixed": [[1, "a", 2.5], ["x"]], + "empty-lists": [[], []], + "dicts": [{"a": 1}, {"b": [1.5, None]}], + "nested-float-lists": [[[0.5, 1.5]], [[2.5]]], + "list-of-datetimes": [[STAMP], [None]], + "bytes-like": [b"ab", b"c"], +} + + +def _rows(values, name="c"): + return [{name: v} for v in values] + + +def _wide_null_heavy(seed=7, count=1000): + """Many rows, ragged keys, about nine values in ten missing or None, three vectors, strings, ints and floats.""" + rng = random.Random(seed) # nosec B311 - fixed test data, not a secret + rows = [] + for i in range(count): + row = {} + if rng.random() < 0.9: + row["i"] = rng.randrange(-1000, 1000) if rng.random() < 0.9 else None + if rng.random() < 0.1: + row["f"] = rng.random() * 10 + if rng.random() < 0.1: + row["s"] = f"s{i}" if rng.random() < 0.8 else None + if rng.random() < 0.1: + row["late-" + str(i % 3)] = i + if rng.random() < 0.1: + row["v"] = [float(np.float32(rng.random())) for _ in range(4)] + if rng.random() < 0.05: + row["m"] = rng.choice([1, 2.5, "x", None, True]) + rows.append(row) + return rows + + +# --- the comparison ------------------------------------------------------------------------------------------------------- + + +def _outcome(fn, *args): + try: + return ("ok", fn(*args)) + except ( + Exception + ) as exc: # noqa: BLE001 - the class of the failure is part of the result + return ("raises", type(exc)) + + +def _same_table(a, b): + assert a.schema == b.schema, (a.schema, b.schema) + assert a.column_names == b.column_names + # not Table.equals: it is False for a NaN next to the same NaN. repr counts NaN, -0.0 and the type of every value + assert repr(a.to_pydict()) == repr(b.to_pydict()) + + +def _same_column(a, b, name): + assert type(a) is type(b), (name, type(a), type(b)) + if isinstance(a, np.ndarray): + assert a.dtype == b.dtype and a.shape == b.shape, (name, a.dtype, b.dtype) + assert repr(a.tolist()) == repr(b.tolist()), name + else: + assert repr(a) == repr(b), name + + +def _check_arrow(rows, columns=None): + want = _outcome(_to_arrow_ref, rows, columns) + have = _outcome(new.to_arrow, rows, columns) + assert have[0] == want[0], (have, want) + if want[0] == "ok": + _same_table(have[1], want[1]) + else: + assert have[1] is want[1] + + +def _check_columns(rows, columns=None): + want = _outcome(_to_columns_ref, rows, columns) + have = _outcome(new.to_columns, rows, columns) + assert have[0] == want[0], (have, want) + if want[0] == "ok": + assert list(have[1]) == list(want[1]) + for name in want[1]: + _same_column(have[1][name], want[1][name], name) + else: + assert have[1] is want[1] + + +@pytest.mark.parametrize("case", sorted(COLUMNS)) +def test_a_single_column_gives_the_same_arrow_table(case): + _check_arrow(_rows(COLUMNS[case])) + + +@pytest.mark.parametrize("case", sorted(COLUMNS)) +def test_a_single_column_gives_the_same_columns(case): + _check_columns(_rows(COLUMNS[case])) + + +def test_a_table_of_all_the_cases_gives_the_same_arrow_table_and_columns(): + names = sorted(COLUMNS) + count = max(len(COLUMNS[n]) for n in names) + rows = [ + {n: COLUMNS[n][i] for n in names if i < len(COLUMNS[n])} for i in range(count) + ] + # a case that raises on its own raises for the table: compare the table without the columns that raise + ok = [n for n in names if _outcome(_to_arrow_ref, _rows(COLUMNS[n]))[0] == "ok"] + rows = [{n: r[n] for n in ok if n in r} for r in rows] + _check_arrow(rows) + ok_columns = [ + n for n in ok if _outcome(_to_columns_ref, _rows(COLUMNS[n]))[0] == "ok" + ] + _check_columns([{n: r[n] for n in ok_columns if n in r} for r in rows]) + + +@pytest.mark.parametrize("seed", [1, 7, 42]) +def test_a_null_heavy_ragged_table_gives_the_same_tables(seed): + rows = _wide_null_heavy(seed) + _check_arrow(rows) + _check_columns(rows) + + +def test_a_chosen_column_list_gives_the_same_tables(): + rows = _wide_null_heavy() + for columns in (["i", "s"], ["v", "nope", "i"], []): + _check_arrow(rows, columns) + _check_columns(rows, columns) + + +@pytest.mark.parametrize("columns", [None, ["a"]]) +def test_no_rows_gives_the_same_tables(columns): + _check_arrow([], columns) + _check_columns([], columns) + + +def test_rows_of_empty_dicts_give_the_same_tables(): + _check_arrow([{}, {}]) + _check_columns([{}, {}]) + + +def test_names_are_the_keys_in_the_order_they_are_first_seen(): + rows = [{"b": 1}, {"a": 1, "b": 2}, {"b": 3}, {"c": 1, "a": 2}, {"d": 1, "b": 1}] + assert new._names(rows, None) == _names_ref(rows, None) == ["b", "a", "c", "d"] + rows = _wide_null_heavy() + assert new._names(rows, None) == _names_ref(rows, None) + assert new._names(rows, ["z", "y"]) == ["z", "y"] + + +def test_the_inference_only_takes_a_column_the_rules_type_the_same(): + """A temporal column and an int above int64 are left to the rules; the rest is the inferred array.""" + assert new._inferred([STAMP, None]) is None + assert new._inferred([date(2026, 1, 2)]) is None + assert new._inferred([2**63]) is None + assert new._inferred([1, None, 3]).type == pa.int64() + assert new._inferred([0.5, 1]).type == pa.float64() + assert new._inferred([True, None]).type == pa.bool_() + assert new._inferred(["a", None]).type == pa.string() + assert new._inferred([None, None]).type == pa.null() + + +def test_to_columns_asks_for_the_kind_of_a_column_once(monkeypatch): + calls = [] + real = new._kind + + def counting(values): + calls.append(len(values)) + return real(values) + + monkeypatch.setattr(new, "_kind", counting) + new.to_columns(_rows([1, None, 3]), None) + assert calls == [3] + calls.clear() + new.to_columns(_rows([[0.5, 1.5], None]), None) + assert calls == [2] + + +# --- the drain ------------------------------------------------------------------------------------------------------------ + +DRAIN_ROWS = 5000 # more than the first reply (256) and two further replies (2048 each) + + +@pytest.fixture +def drain_db(temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Drain") + temp_db.command("sql", "CREATE PROPERTY Drain.id INTEGER") + temp_db.insert_many( + "Drain", + [{"id": i, "txt": f"t{i}", "val": i / 8.0} for i in range(DRAIN_ROWS)], + ) + return temp_db + + +SELECT_DRAIN = "SELECT id, txt, val FROM Drain ORDER BY id" + + +def test_the_drain_returns_the_rows_row_by_row_iteration_returns(drain_db): + by_rows = [r.to_dict() for r in drain_db.query("sql", SELECT_DRAIN)] + drained = drain_db.query("sql", SELECT_DRAIN)._drain_dicts() + assert len(drained) == DRAIN_ROWS + assert drained == by_rows + assert drained[0] == {"id": 0, "txt": "t0", "val": 0.0} + assert drained[-1] == { + "id": DRAIN_ROWS - 1, + "txt": f"t{DRAIN_ROWS - 1}", + "val": (DRAIN_ROWS - 1) / 8.0, + } + + +@pytest.mark.parametrize("taken", [1, 255, 256, 257, 2304, 4999]) +def test_the_drain_after_some_rows_were_taken_returns_the_rest(drain_db, taken): + rs = drain_db.query("sql", SELECT_DRAIN) + head = [next(rs).to_dict() for _ in range(taken)] + rest = rs._drain_dicts() + want = [r.to_dict() for r in drain_db.query("sql", SELECT_DRAIN)] + assert head + rest == want + + +def test_the_drain_of_an_empty_result_and_a_second_drain_are_empty(drain_db): + rs = drain_db.query("sql", "SELECT id FROM Drain WHERE id < 0") + assert rs._drain_dicts() == [] + rs = drain_db.query("sql", SELECT_DRAIN) + assert len(rs._drain_dicts()) == DRAIN_ROWS + assert rs._drain_dicts() == [] + + +def test_to_arrow_and_to_columns_of_a_result_equal_the_reference_over_its_rows( + drain_db, +): + rows = drain_db.query("sql", SELECT_DRAIN)._drain_dicts() + table = drain_db.query("sql", SELECT_DRAIN).to_arrow() + _same_table(table, _to_arrow_ref(rows)) + columns = drain_db.query("sql", SELECT_DRAIN).to_columns() + want = _to_columns_ref(rows) + for name in want: + _same_column(columns[name], want[name], name)