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