diff --git a/docs/modules/ROOT/pages/how-to/vectors.adoc b/docs/modules/ROOT/pages/how-to/vectors.adoc index 84beb2db..4396b194 100644 --- a/docs/modules/ROOT/pages/how-to/vectors.adoc +++ b/docs/modules/ROOT/pages/how-to/vectors.adoc @@ -110,9 +110,9 @@ print(db.count_type("Chunk")) 200 ---- -`GraphBatch.create_vertices` rows and the properties of `new_edge` and `new_edges` take the same arrays. On the `arcadedb` (native) package, a non-empty one-dimensional `float32` array in a row or an edge property crosses in a binary side buffer, and the stored values are bit-identical to the ones the text path gave. A list or a `float64` array still crosses as number text, so make the array with `np.asarray(v, dtype=np.float32)` before the load. A laptop A/B on the native wheel (100,000 vertices for each load, `ab.py --noise`, ABBA order, 12 clean pairs for each figure, clean A/A control) gave this wall time of the new path against the old one (95 % interval of the median pair ratio): 0.517 [0.515, 0.520] at 128 dimensions, 0.218 [0.217, 0.219] at 96, and 0.271 [0.270, 0.272] at 64. A load without vectors was unchanged (1.009 [1.003, 1.012]). +`GraphBatch.create_vertices` rows and the properties of `new_edge` and `new_edges` take the same arrays. On the `arcadedb` (native) package, a non-empty one-dimensional `float32` array in a row or an edge property crosses in a binary side buffer, and the stored values are bit-identical to the ones the text path gave. A list or a `float64` array still crosses as number text, so make the array with `np.asarray(v, dtype=np.float32)` before the load. When a property holds a `float32` array of one length in every row of a `create_vertices` call, the whole column crosses as one float32 matrix in that buffer, with no tag or Python call for each row. A laptop A/B on the native wheel (100,000 vertices for each load, `ab.py --noise`, ABBA order, 12 clean pairs for each figure, clean A/A control) gave this wall time of the new path against the old one (95 % interval of the median pair ratio): 0.517 [0.515, 0.520] at 128 dimensions, 0.218 [0.217, 0.219] at 96, and 0.271 [0.270, 0.272] at 64. A load without vectors was unchanged (1.009 [1.003, 1.012]). -On the `arcadedb-jvm` package, `create_vertices` already sends the values of a property declared `ARRAY_OF_FLOATS` as one binary float32 matrix when every row holds a list or array of numbers of one length. So `float32` arrays in a property declared `ARRAY_OF_FLOATS` take the binary path on both packages. +On the `arcadedb-jvm` package, `create_vertices` sends the values of a property declared `ARRAY_OF_FLOATS` as one binary float32 matrix when every row holds a list, a numpy array or a Java `float[]` or `double[]` of numbers of one length. A property that is not declared takes the same matrix when every row holds a `float32` numpy array or a Java `float[]` (what `to_java_float_array` returns on this package) of one length, and the stored value is a `float[]`, as before. A list or a `float64` array on an undeclared property keeps its doubles and the older path. So `float32` arrays and `to_java_float_array` values take the binary path on both packages, declared or not. == Search diff --git a/native/src/main/java/org/humemai/nativelib/NativeExtras.java b/native/src/main/java/org/humemai/nativelib/NativeExtras.java index a62574ad..f359fe80 100644 --- a/native/src/main/java/org/humemai/nativelib/NativeExtras.java +++ b/native/src/main/java/org/humemai/nativelib/NativeExtras.java @@ -73,6 +73,7 @@ import org.graalvm.word.WordFactory; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -250,6 +251,45 @@ static Object[] varargs(final Object props) { return out; } + /** + * The float32 vector columns of create_vertices rows: vecs is {"names": [..], "dims": [..], "offs": [..]} and the side buffer holds, for + * column c, one vector after the other (rows * dims[c] floats) from float offset offs[c]. Row i gets the property names[c] as a float[] + * of its slice, appended to its property array: the float[] that a {"\u0001":["B",offset,count]} tag per row gave, from one matrix and + * no tag per row. A null or empty vecs leaves the rows as they are. + */ + private static void withVectors(final List flats, final Object vecs, final long buf, final long bufFloats) { + if (!(vecs instanceof Map v)) + return; + final List names = (List) v.get("names"), dims = (List) v.get("dims"), offs = (List) v.get("offs"); + final int columns = names.size(), rows = flats.size(); + if (columns == 0) + return; + if (buf == 0 || dims.size() != columns || offs.size() != columns) + throw new IllegalArgumentException("vector columns without a matching side buffer"); + final int[] dim = new int[columns]; + final long[] off = new long[columns]; + for (int c = 0; c < columns; c++) { + dim[c] = ((Number) dims.get(c)).intValue(); + off[c] = ((Number) offs.get(c)).longValue(); + if (dim[c] < 1 || off[c] < 0 || off[c] + (long) rows * dim[c] > bufFloats) + throw new IllegalArgumentException("vector column outside the side buffer"); + } + final CFloatPointer fp = WordFactory.pointer(buf); + for (int i = 0; i < rows; i++) { + final Object[] row = flats.get(i); + final Object[] grown = Arrays.copyOf(row, row.length + 2 * columns); + for (int c = 0; c < columns; c++) { + final float[] f = new float[dim[c]]; + final int base = (int) (off[c] + (long) i * dim[c]); + for (int k = 0; k < dim[c]; k++) + f[k] = fp.read(base + k); + grown[row.length + 2 * c] = names.get(c); + grown[row.length + 2 * c + 1] = f; + } + flats.set(i, grown); + } + } + static String ridList(final RID[] rids) { final StringBuilder sb = new StringBuilder("["); for (int i = 0; i < rids.length; i++) { @@ -488,8 +528,10 @@ else if (opName.equals("create_vertices")) final RID[] rids; if (!(a.get("rows") instanceof Integer)) rids = batch.createVertices((String) a.get("type"), intOf(a.get("count"), 0)); - else + else { + withVectors(flats, a.get("vecs"), buf, bufFloats); rids = batch.createVertices((String) a.get("type"), flats.toArray(new Object[0][])); + } return ArcadeNative.out("{\"ok\":true,\"rids\":" + ridList(rids) + "}"); } case "new_edge" -> batch.newEdge(new RID((String) a.get("src")), (String) a.get("type"), new RID((String) a.get("dst")), diff --git a/src/arcadedb/_native/codec.py b/src/arcadedb/_native/codec.py index cf6d3f15..51d5d74e 100644 --- a/src/arcadedb/_native/codec.py +++ b/src/arcadedb/_native/codec.py @@ -1,5 +1,7 @@ """Python <-> native JSON: query parameters out, typed rows in.""" +import ctypes +import itertools import json import os from collections.abc import Mapping @@ -276,6 +278,64 @@ def encode_gb_args(args): return text, buf, state[1] +_STR_TYPE = frozenset((str,)) + + +def _plain_rows(rows): + """True when every row is a dict whose keys are str and whose values are plain scalars: :func:`_prep` and the default hook would + change nothing, so ``json.dumps`` writes the rows as they are. The census runs in C (a set of types per level), not in Python per row. + """ + if not set(map(type, set(itertools.chain.from_iterable(rows)))) <= _STR_TYPE: + return False + return ( + set(map(type, itertools.chain.from_iterable(map(dict.values, rows)))) <= _PLAIN + ) + + +def _side_buffer(flat): + """The float32 array `flat` as the side-buffer argument of ``arcadedb_n_gb_bin``: a ctypes view of its memory (no copy), or its bytes + when the array is read-only or not contiguous.""" + try: + return (ctypes.c_char * flat.nbytes).from_buffer(flat) + except (TypeError, ValueError): + return flat.tobytes() + + +def encode_vertex_args(type_name, rows, names, matrices): + """The ``args`` of ``arcadedb_n_gb_bin`` for ``create_vertices`` whose vectors are float32 matrices: ``(text, side buffer, float count)``. + + `rows` hold the properties other than `names`; matrices[c] is the ``(len(rows), dim)`` float32 array of the property names[c], one + row of it per vertex. The matrices are the front of the side buffer, back to back, and the text carries their names, dimensions and + float offsets (``"vecs"``): no tag and no Python call per row. A value of `rows` that needs the hook (a datetime, a float32 array that + did not qualify, a list of floats) is written as in :func:`encode_gb_args`, after the matrices in the buffer. + """ + offs, total = [], 0 + for matrix in matrices: + offs.append(total) + total += matrix.size + args = { + "type": type_name, + "rows": rows, + "vecs": { + "names": names, + "dims": [int(m.shape[1]) for m in matrices], + "offs": offs, + }, + } + if _plain_rows(rows): + text, arrays = _dumps(args), list(matrices) + else: + state = [list(matrices), total] + text = json.dumps(_prep(args), default=lambda v: _gb_default(v, state)) + arrays = state[0] + flat = ( + arrays[0].reshape(-1) + if len(arrays) == 1 + else _np.concatenate([a.reshape(-1) for a in arrays]) + ) + return text, _side_buffer(flat), flat.size + + _INT64_MAX = 2**63 - 1 _NS_DATETIME = None if _np is None else _np.dtype("datetime64[ns]") diff --git a/src/arcadedb/_native/extras.py b/src/arcadedb/_native/extras.py index 5f8a4702..a5ff9eca 100644 --- a/src/arcadedb/_native/extras.py +++ b/src/arcadedb/_native/extras.py @@ -19,8 +19,18 @@ from ..callbacks import check_event, check_function from ..exceptions import ArcadeDBError from ..graph_batch import GraphBatch as _BaseGraphBatch +from ..type_conversion import take_vector_columns from . import interrupts, lib, not_yet -from .codec import _PLAIN, _default, _np, _prep, encode_gb_args, loads, sample_array +from .codec import ( + _PLAIN, + _default, + _np, + _prep, + encode_gb_args, + encode_vertex_args, + loads, + sample_array, +) # -------------------------------------------------------------------------------------------------------- upcall @@ -28,6 +38,7 @@ _UPCALL_TYPE = ctypes.CFUNCTYPE( None, ctypes.c_longlong, ctypes.c_char_p, ctypes.c_void_p ) +_DICT_OR_NONE = frozenset((dict, type(None))) _callbacks: dict = {} _callback_ids = itertools.count(1) _callbacks_lock = threading.Lock() @@ -180,6 +191,30 @@ def _op(self, prefix: str, op: str, **args) -> dict: return _call(prefix, "arcadedb_n_gb", self._gb, lib.b(op), lib.b(text)) return _call(prefix, "arcadedb_n_gb", self._gb, lib.b(op), lib.b(_enc(args))) + def _create_chunk(self, prefix: str, type_name: str, rows: list) -> list: + """One crossing and one transaction for `rows` (dicts): the RIDs of the new vertices. + + A property whose value in every row is a float32 numpy array of one length crosses as one float32 matrix in the side buffer + (``take_vector_columns``, ``encode_vertex_args``), with no tag, hook call or Python code per row; every other value is + written as :func:`encode_gb_args` writes it. + """ + taken = take_vector_columns(rows, (), undeclared=True) if lib.GB_BIN else None + if taken is None: + return self._op(prefix, "create_vertices", type=type_name, rows=rows)[ + "rids" + ] + names, matrices, rest = taken + text, buf, count = encode_vertex_args(type_name, rest, names, matrices) + return _call( + prefix, + "arcadedb_n_gb_bin", + self._gb, + lib.b("create_vertices"), + lib.b(text), + buf, + count, + )["rids"] + def new_vertex(self, type_name: str): from .records import Vertex @@ -229,27 +264,14 @@ def _to_java_varargs(self, properties: dict) -> dict: raise TypeError(f"property names must be str, got {key!r}") return properties - def _create_vertices_json_bulk(self, type_name: str, rows): - """The bulk path (JSON rows in, RIDs out): the RIDs, or None for rows JSON would change (the caller falls back).""" - from ..type_conversion import json_bulk_scalar_ok - - for row in rows: - for key, value in (row or {}).items(): - if type(key) is not str or not json_bulk_scalar_ok(value): - return None - return self._op( - f"Failed to create batch vertices for type '{type_name}'", - "create_vertices", - type=type_name, - rows=[r or {} for r in rows], - )["rids"] - def create_vertices(self, type_name: str, count_or_properties) -> list: """The RIDs of the new vertices. Rows go in chunks of the batch's ``commit_every`` (1000 unless graph_batch() got another value; 0: 10,000), one crossing and one transaction per chunk: chunks that already went in stay stored if a later chunk fails. - A non-empty one-dimensional ``float32`` numpy array in a row crosses in a binary side buffer; a list or a ``float64`` array - crosses as number text (``new_edge`` and ``new_edges`` props follow the same rule). + A property whose value in every row is a non-empty one-dimensional ``float32`` numpy array of one length crosses as one + float32 matrix in a binary side buffer, one buffer per property and chunk; any other ``float32`` array of a row crosses in the side + buffer one array at a time, and a list or a ``float64`` array as number text (``new_edge`` and ``new_edges`` props follow the + one-array rule). """ self._check_not_closed() db = self._db @@ -274,23 +296,24 @@ def create_vertices(self, type_name: str, count_or_properties) -> list: ] return rids rows = list(count_or_properties) - for row in rows: - if row is not None and not isinstance(row, dict): - raise TypeError( - "create_vertices() property rows must be dicts or None" - ) + kinds = set(map(type, rows)) + if ( + not kinds <= _DICT_OR_NONE + ): # exact types in C; only an odd row pays the isinstance loop + for row in rows: + if row is not None and not isinstance(row, dict): + raise TypeError( + "create_vertices() property rows must be dicts or None" + ) # chunked, as the JPype wrapper does: one crossing, one JSON text and one reply per chunk, so a 1M-vertex load # never holds the text, the parsed rows and the reply of all of it at once (the Serial GC copies what lives) rids = [] + has_none = type(None) in kinds for start in range(0, len(rows), chunk): - rids.extend( - self._op( - prefix, - "create_vertices", - type=type_name, - rows=[r or {} for r in rows[start : start + chunk]], - )["rids"] - ) + part = rows[start : start + chunk] + if has_none: + part = [r or {} for r in part] + rids.extend(self._create_chunk(prefix, type_name, part)) if not rows: rids = self._op(prefix, "create_vertices", type=type_name, rows=[])[ "rids" diff --git a/src/arcadedb/graph_batch.py b/src/arcadedb/graph_batch.py index 03495844..094e1d23 100644 --- a/src/arcadedb/graph_batch.py +++ b/src/arcadedb/graph_batch.py @@ -12,9 +12,11 @@ from .exceptions import ArcadeDBError from .graph import Document, Vertex from .type_conversion import ( + classify_vertex_rows, convert_python_to_java, declared_float_vectors, - is_vector_candidate, + declared_scalars_ok, + json_bulk_dumps, json_bulk_dumps_tagged, json_bulk_scalar_ok, json_bulk_shape_ok, @@ -190,9 +192,12 @@ def create_vertices( count_or_properties: Either an integer vertex count or an iterable of per-vertex property dictionaries. Use `None` or `{}` for vertices without properties when passing an iterable. A property declared - ``ARRAY_OF_FLOATS`` whose value in every row is a list or a numpy - array of numbers of one length crosses as one binary float32 - matrix instead of JSON text. + ``ARRAY_OF_FLOATS`` whose value in every row is a list, a numpy + array or a Java ``float[]`` or ``double[]`` of numbers of one length + crosses as one binary float32 matrix instead of JSON text, and so + does any other property whose value in every row is a float32 + numpy array or a Java ``float[]`` (``to_java_float_array``) of one + length: one buffer per property and chunk, stored as a ``float[]``. A list (not one of only numbers), tuple, dict, date-time, date, Decimal or bytes value also stays on the bulk path: it crosses inside the JSON as a tagged value and is stored as the @@ -250,64 +255,49 @@ def create_vertices( def _create_vertices_json_bulk(self, type_name, rows): """Bulk path: rows cross as one JSON string, RIDs return as one joined string (two bulk copies instead of per-value marshaling, measured - ~2.4x). The float vectors of ARRAY_OF_FLOATS properties cross as one - binary matrix beside the JSON instead of as text (split_float_vectors). + ~2.4x). The float vectors of ARRAY_OF_FLOATS properties, and the float32 + arrays and Java float[] of any property, cross as one binary matrix + beside the JSON instead of as text (split_float_vectors). A list (not one of only numbers), tuple, dict, date-time, date, Decimal or bytes value crosses inside the JSON as the tagged value json_bulk_dumps_tagged writes and VertexBatcher decodes (the stored value is what the property-matrix path stores). An input whose values - are all scalars takes the untagged json.dumps path as before. A chunk - holding a value the text cannot carry (a NaN, an integer beyond 64 - bits or a date-time inside a list, a subnormal) goes through the - property-matrix path on its own. Returns None when - unavailable/unsuitable so the caller falls back to the property-matrix - path for the whole input.""" + are all scalars is written by json_bulk_dumps. The values are checked + in the one json.dumps pass of each chunk (see type_conversion._dumps_once), + not in a loop over the rows before it: a chunk holding a value the text + cannot carry (a NaN, an integer beyond 64 bits, a lone surrogate, a + subnormal, or a date-time inside a list) goes through the property-matrix + path on its own. Returns None when unavailable/unsuitable (a property + name that is not a str, a type the text does not carry) so the caller + falls back to the property-matrix path for the whole input.""" try: vertex_batcher = jpype.JClass("com.arcadedb.python.VertexBatcher") declared = set(declared_float_vectors(self._java_db, type_name)) except Exception: return None - tagged = False # a value that needs the tagged text; False keeps the untagged code of before - for row in rows: - if row: - for key, value in row.items(): - if type(key) is not str: - # a property name that is not a str, which json.dumps would - # write as text ({1: "a"} as "1") where the matrix path - # refuses it - return None - if key in declared and is_vector_candidate(value): - continue # the vector split below takes it, or its chunk goes the matrix way - if not json_bulk_scalar_ok(value): - if not json_bulk_shape_ok(value): - # a type the tagged JSON does not carry (a set, a numpy - # or pandas value, any subclass, a list of only numbers: - # the matrix path converts those), or a scalar JSON - # would change: an integer beyond 64 bits, NaN or - # Infinity, a lone surrogate - return None - tagged = True - - import json + tagged = classify_vertex_rows(rows, declared) + if tagged is None: + # a property name that is not a str (json.dumps would write 1 as "1" where the matrix path refuses it), or a type the + # tagged JSON does not carry (a set, a numpy or pandas value, any subclass, a list of only numbers: the matrix path + # converts those) + return None + dumps = json_bulk_dumps_tagged if tagged else json_bulk_dumps # chunked so huge ingests never materialize one giant JSON string rids = [] for start in range(0, len(rows), self._BULK_CHUNK): chunk = rows[start : start + self._BULK_CHUNK] - split = split_float_vectors(self._java_db, type_name, chunk, declared) + split = split_float_vectors( + self._java_db, type_name, chunk, declared, undeclared=True + ) rest = split[0] if split else chunk - if not tagged: - text = None - if not declared or self._json_bulk_ok(rest): - text = json.dumps([row or {} for row in rest]) - else: - text = None - if self._json_bulk_declared_ok(rest, declared): - try: - text = json_bulk_dumps_tagged([row or {} for row in rest]) - except (TypeError, ValueError): - text = None # a value the text cannot carry, see the docstring + text = None + if declared_scalars_ok(rest, declared): + try: + text = dumps([row or {} for row in rest]) + except (TypeError, ValueError): + text = None # a value the text cannot carry, see the docstring if text is None: # a list the matrix cannot carry (ragged, a null, not numbers) # is left in this chunk, or a value the text would change: the @@ -330,25 +320,6 @@ def _create_vertices_json_bulk(self, type_name, rows): rids.extend(joined.split(";")) return rids - @staticmethod - def _json_bulk_ok(rows): - for row in rows: - if row: - for value in row.values(): - if not json_bulk_scalar_ok(value): - return False - return True - - @staticmethod - def _json_bulk_declared_ok(rows, declared): - """False when a declared ARRAY_OF_FLOATS property still holds a list (or any non-scalar) after the vector split.""" - for row in rows: - if row: - for key in declared: - if key in row and not json_bulk_scalar_ok(row[key]): - return False - return True - def new_edge( self, source_vertex_or_rid, diff --git a/src/arcadedb/type_conversion.py b/src/arcadedb/type_conversion.py index 1eca63ed..cc040da8 100644 --- a/src/arcadedb/type_conversion.py +++ b/src/arcadedb/type_conversion.py @@ -5,8 +5,10 @@ developer experience and integration with Python ecosystem (pandas, numpy, etc.). """ +import itertools import json import math +import operator import re import sys from datetime import date, datetime, time, timedelta, timezone @@ -804,16 +806,38 @@ def _dumps_checked(rows: Any, **kwargs: Any) -> str: return text -def json_bulk_dumps(rows: Any) -> str: - """json.dumps for the bulk paths: raises ValueError (or TypeError) when - any value in `rows` would change on the way to the engine. +def _int64_overflow_pattern() -> str: + """The digit runs of 2**63 and above, which a JSON integer token beyond 64 bits has. + + A run of 20 digits or more, or of 19 digits that is not below 9223372036854775808. A run that follows a digit or a "." is not an + integer token (the digits after the point of a float such as 0.00012345678901234568 are no token), so only a run start counts. + """ + limit = str(2**63) + parts = [r"[0-9]{20,}"] + for i, digit in enumerate(limit): + low = int(digit) if i == len(limit) - 1 else int(digit) + 1 + if low <= 9: + parts.append(f"{limit[:i]}[{low}-9][0-9]{{{len(limit) - 1 - i}}}") + return r"(? None: + """The per-row walk: raise ValueError (or TypeError) for the first row, in row order, that holds a value the JSON text would change. + + This is the exact reference. The bulk dumps do not run it on every call: `_dumps_once` runs it only when the text or the types of + the rows show that a value may be bad, and when json.dumps itself refused the rows, so the error is always the walk's. + `tagged` is insert_many's format (see json_bulk_dumps_tagged), where the kinds bulk_tag_default writes are allowed in a row. """ int_min, int_max = _INT64_MIN, _INT64_MAX + kinds = _TAGGED_KINDS for row in rows: for key in row: if type(key) is not str: @@ -832,11 +856,87 @@ def json_bulk_dumps(rows: Any) -> str: ) elif kind is float or kind is bool or value is None: continue + elif tagged and kind in kinds: + continue # bulk_tag_default checks it as json.dumps writes it + elif tagged and kind is dict and _TAG in value: + raise ValueError( + "a map keyed by the tag marker cannot go through the JSON bulk path" + ) + elif tagged and ( + isinstance(value, (datetime, date, Decimal, bytes, bytearray)) + or (_np is not None and isinstance(value, _np.generic)) + or _is_pandas_null(kind) + ): + continue # a subclass (pandas.Timestamp), a numpy scalar, pandas NA or NaT else: _json_bulk_check( value ) # nested lists and dicts, and subclasses of the types above - return _dumps_checked(rows) + + +def _kind_is_tagged(kind: type) -> bool: + """Whether bulk_tag_default (not the walk) checks a value of this type: the kinds `_check_rows` lets through in tagged mode.""" + return ( + kind in _TAGGED_KINDS + or issubclass(kind, (datetime, date, Decimal, bytes, bytearray)) + or (_np is not None and issubclass(kind, _np.generic)) + or _is_pandas_null(kind) + ) + + +def _rows_need_walk(rows: Any, tagged: bool) -> bool: + """Whether the types of the rows leave a check that the JSON text cannot show: only exact dict rows with str keys and values of + the exact scalar types (and, tagged, the kinds bulk_tag_default checks) are decided by the text. The census runs in C: one pass + over the keys and one over the values, no Python-level loop per row. + """ + if not set(map(type, rows)) <= _DICT_ONLY: + return True + if not set(map(type, itertools.chain.from_iterable(rows))) <= _STR_ONLY: + return True + extra = ( + set(map(type, itertools.chain.from_iterable(map(dict.values, rows)))) + - _BULK_SCALAR_TYPES + ) + if not extra: + return False + return not (tagged and all(map(_kind_is_tagged, extra))) + + +def _dumps_once(rows: Any, tagged: bool) -> str: + """One json.dumps pass for the bulk paths, with the checks of `_check_rows` decided from what that pass shows. + + A str or int the text would change cannot be seen by a json.dumps hook (the encoder writes both itself), so the text decides: + a surrogate escape (a lone surrogate) or an integer token of 2**63 and above sends the rows to the walk, and so do rows whose + types the text cannot vouch for (`_rows_need_walk`). When json.dumps itself refuses the rows (NaN, an unknown type, a + subnormal, a hook error) the walk runs first, as the validation loop used to run before json.dumps, so the error of the first + bad row in row order is still the walk's; if the walk finds nothing, json.dumps's own error stands. A false alarm costs the walk + and changes nothing else. + """ + try: + text = ( + _dumps_checked(rows, default=bulk_tag_default) + if tagged + else _dumps_checked(rows) + ) + except (TypeError, ValueError): + _check_rows(rows, tagged) + raise + if ( + _rows_need_walk(rows, tagged) + or _SURROGATE_ESCAPE.search(text) + or _INT64_OVERFLOW.search(text) + ): + _check_rows(rows, tagged) + return text + + +def json_bulk_dumps(rows: Any) -> str: + """json.dumps for the bulk paths: raises ValueError (or TypeError) when + any value in `rows` would change on the way to the engine. + + `rows` is a list of dicts (one per record). The rows are written once; see `_dumps_once` for how the values are checked. + """ + return _dumps_once(rows, False) # --------------------------------------------------------------------------- @@ -933,43 +1033,7 @@ def json_bulk_dumps_tagged(rows: Any) -> str: through, written by bulk_tag_default. Raises TypeError or ValueError exactly where json_bulk_dumps does for every other value (a datetime inside a list or map included), so the caller can send that row down the per-row path. """ - int_min, int_max = _INT64_MIN, _INT64_MAX - tagged = _TAGGED_KINDS - for row in rows: - for key in row: - if type(key) is not str: - raise ValueError("a non-str key cannot go through the JSON bulk path") - for value in row.values(): - kind = type(value) - if kind is str: - if not value.isascii() and not _str_survives_json(value): - raise ValueError( - "a string with a lone surrogate cannot go through the JSON bulk path" - ) - elif kind is int: - if not int_min <= value <= int_max: - raise ValueError( - f"an integer beyond 64 bits cannot go through the JSON bulk path: {value}" - ) - elif kind is float or kind is bool or value is None: - continue - elif kind in tagged: - continue # bulk_tag_default checks it as json.dumps writes it - elif kind is dict and _TAG in value: - raise ValueError( - "a map keyed by the tag marker cannot go through the JSON bulk path" - ) - elif ( - isinstance(value, (datetime, date, Decimal, bytes, bytearray)) - or (_np is not None and isinstance(value, _np.generic)) - or _is_pandas_null(kind) - ): - continue # a subclass (pandas.Timestamp), a numpy scalar, pandas NA or NaT - else: - _json_bulk_check( - value - ) # nested lists and dicts, and subclasses of the types above - return _dumps_checked(rows, default=bulk_tag_default) + return _dumps_once(rows, True) def bulk_row_ok(row: Any, skip=()) -> bool: @@ -983,6 +1047,80 @@ def bulk_row_ok(row: Any, skip=()) -> bool: return True +_ROWS_OR_NONE = frozenset((dict, type(None))) + + +def _column(rows: Any, key: str): + """The value of `key` in every non-empty row (None where the row lacks it), as one C-level iterator.""" + return map(dict.get, filter(None, rows), itertools.repeat(key)) + + +def classify_vertex_rows(rows: Any, declared: Any): + """What GraphBatch.create_vertices must know about its rows before it writes any chunk, decided by types. + + Returns None when the rows must all go the property-matrix way (a property name that is not a str, or a value of a type the + tagged JSON does not carry: a set, a numpy or pandas value, a list or tuple of only numbers, any other object); otherwise True when + some value needs the tagged text (json_bulk_dumps_tagged) and False when json_bulk_dumps carries every row. A value of an exact + scalar type is not judged here: whether the text carries it (a NaN, an integer beyond 64 bits, a lone surrogate) is decided in + the dumps of its chunk. The property of `declared` (ARRAY_OF_FLOATS) that holds a list or a numpy array is left to + split_float_vectors. The census over the keys and the values runs in C; a Python-level loop runs only over a column that + holds a type other than the exact scalars. + """ + if not set(map(type, rows)) <= _ROWS_OR_NONE: + return _classify_walk( + rows, declared + ) # a row that is no dict: the walk raises what it always raised + keys = set(itertools.chain.from_iterable(filter(None, rows))) + if not set(map(type, keys)) <= _STR_ONLY: + return None + tagged = False + for key in keys: + odd = set(map(type, _column(rows, key))) - _BULK_SCALAR_TYPES + if not odd: + continue + if key in declared and all(_is_vector_kind(kind) for kind in odd): + continue # the vectors of a declared column + if all(_is_float_array_kind(kind) for kind in odd): + continue # a column of arrays: take_vector_columns takes the float32 ones; a chunk it leaves goes the matrix way + for value in _column(rows, key): + if type(value) in _BULK_SCALAR_TYPES or ( + key in declared and is_vector_candidate(value) + ): + continue + if not json_bulk_scalar_ok(value): + if not json_bulk_shape_ok(value): + return None + tagged = True + return tagged + + +def _classify_walk(rows: Any, declared: Any): + tagged = False + for row in rows: + if row: + for key, value in row.items(): + if type(key) is not str: + return None + if key in declared and is_vector_candidate(value): + continue + if not json_bulk_scalar_ok(value): + if not json_bulk_shape_ok(value): + return None + tagged = True + return tagged + + +def declared_scalars_ok(rows: Any, declared: Any) -> bool: + """False when a declared ARRAY_OF_FLOATS property still holds a list (or any non-scalar) after the vector split.""" + for key in declared: + if set(map(type, _column(rows, key))) <= _BULK_SCALAR_TYPES: + continue # scalars: the dumps of the chunk decides whether the text carries them + for row in rows: + if row and key in row and not json_bulk_scalar_ok(row[key]): + return False + return True + + _FLOAT_VECTORS = None # the bridge's FloatVectors class once resolved; False when the bridge jar lacks it @@ -1004,41 +1142,103 @@ def declared_float_vectors(java_db, type_name): return [str(name) for name in vectors.declared(java_db, type_name)] +_JAVA_VECTOR_ARRAYS = None # (float[] class, double[] class) of JPype once the JVM runs; () where there is no JVM (the native wheel) + + +def _java_vector_arrays(): + """The JPype classes of a Java ``float[]`` and a ``double[]`` (what ``to_java_float_array`` returns on the JVM wheel), () without a JVM. + + Resolved once the JVM has started: JPype builds an array class from a Java class, so it cannot be done at import. + """ + global _JAVA_VECTOR_ARRAYS + if _JAVA_VECTOR_ARRAYS is None: + try: + if not jpype.isJVMStarted(): + return () + _JAVA_VECTOR_ARRAYS = ( + jpype.JArray(jpype.JFloat), + jpype.JArray(jpype.JDouble), + ) + except Exception: # the native wheel's stand-in for jpype has no such classes + _JAVA_VECTOR_ARRAYS = () + return _JAVA_VECTOR_ARRAYS + + +def _is_vector_kind(kind): + """The type of a value that a declared ARRAY_OF_FLOATS column may hold as a vector: a list, a numpy array, a Java float[] or double[].""" + return ( + kind is list + or (_np is not None and issubclass(kind, _np.ndarray)) + or issubclass(kind, _java_vector_arrays()) + ) + + +def _is_float_array_kind(kind): + """The type of a value that may be a float32 vector on any property: a numpy array or a Java float[].""" + arrays = _java_vector_arrays() + return (_np is not None and issubclass(kind, _np.ndarray)) or ( + bool(arrays) and issubclass(kind, arrays[0]) + ) + + def is_vector_candidate(value): - """A value split_float_vectors may take out of the JSON: a list or a numpy array.""" - return type(value) is list or (_np is not None and isinstance(value, _np.ndarray)) + """A value split_float_vectors may take out of the JSON: a list, a numpy array, or a Java float[] or double[].""" + return _is_vector_kind(type(value)) -def split_float_vectors(java_db, type_name, rows, declared=None): - """Take the float vectors of `rows` out of the JSON text of a bulk load. +def _is_float32_vector(value): + """A numpy float32 array or a Java float[]: the value the property-matrix path stores as a float[] on any property.""" + if _np is not None and isinstance(value, _np.ndarray): + return value.dtype == _np.float32 + arrays = _java_vector_arrays() + return bool(arrays) and isinstance(value, arrays[0]) - A list of floats costs Python about 450 ns a number to write as text and the - engine about as much again to parse it, so a vector load spent most of its - time on the text. For each property of `type_name` that is declared - ARRAY_OF_FLOATS, whose value in EVERY row is a list or a numpy array of - numbers of one length, the vectors cross as one float32 matrix instead - (``FloatVectors.Matrix``, one buffer copy). The engine converts the numbers of a list to float for - such a property, so the float32 matrix stores the same values. Every other - property, and a column that is missing in a row, ragged, not numeric, or - null in a row, stays in the JSON. - Returns ``(rest_rows, names, matrices)``: `rest_rows` are the rows without - those properties, `names` a ``String[]`` and `matrices` an ``Object[]`` of - ``FloatVectors.Matrix`` for the bridge; or None when no property qualifies - (the caller sends `rows` as they are). +def _without_keys(rows, gone): + """`rows` (each one holds every name of `gone`) as new dicts without those keys; the values are the same objects. + + Rows that all have the same keys in the same order are rebuilt from columns by C iterators (itemgetter, zip, dict), with no Python + code per row; any other rows take the dict comprehension. + """ + if len(set(map(tuple, rows))) == 1: + keep = [key for key in rows[0] if key not in gone] + if not keep: + return list(map(dict, itertools.repeat((), len(rows)))) + columns = [list(map(operator.itemgetter(key), rows)) for key in keep] + return list(map(dict, map(zip, itertools.repeat(keep), zip(*columns)))) + gone = set(gone) + return [{k: v for k, v in row.items() if k not in gone} for row in rows] + + +def take_vector_columns(rows, declared, undeclared=False): + """The float vector columns of `rows` as float32 matrices, and the rows without them (the one split of both wheels). + + A declared (ARRAY_OF_FLOATS) property qualifies when its value in EVERY row is a list, a numpy array or a Java float[] or double[] of + numbers of one length: the engine converts the numbers to float for such a property, so a float32 matrix stores the same values. + With `undeclared`, any other property qualifies when its value in every row is a float32 numpy array or a Java float[] of one length + (the first row decides which names are tried, the stacked column must come out as float32): the property-matrix path stores such a + value as a float[], and so does the matrix. A float64 array, a list, a Java double[] on an undeclared property, a missing value, a + null, a ragged or non-numeric column does not qualify and stays in the rows. + + Returns None when no column qualifies, else ``(names, matrices, rest)``: the matrices are C-contiguous float32 arrays of shape + ``(len(rows), dim)``, `rest` the rows without those properties. """ if _np is None or not rows: return None - vectors = _float_vectors_class() - if vectors is None: - return None - if declared is None: - declared = declared_float_vectors(java_db, type_name) + declared = set(declared) + names = list(declared) + first = rows[0] + if undeclared and isinstance(first, dict): + names.extend( + key + for key, value in first.items() + if key not in declared and _is_float32_vector(value) + ) taken = [] matrices = [] - for name in declared: + for name in names: try: - column = [row[name] for row in rows] + column = list(map(operator.itemgetter(name), rows)) except (KeyError, TypeError): continue if not is_vector_candidate(column[0]): @@ -1049,24 +1249,53 @@ def split_float_vectors(java_db, type_name, rows, declared=None): continue if array.ndim != 2 or array.shape[1] == 0 or array.dtype.kind not in "fiu": continue + if name not in declared and array.dtype != _np.float32: + continue with _np.errstate( all="ignore" ): # a double beyond float32 becomes Infinity, as in the engine - flat = _np.ascontiguousarray(array, dtype=_np.float32).reshape(-1) + matrix = _np.ascontiguousarray(array, dtype=_np.float32) taken.append(name) - matrices.append( - vectors.Matrix(jpype.JArray(jpype.JFloat)(flat), int(array.shape[1])) - ) + matrices.append(matrix) if not taken: return None - if len(taken) == 1: - only = taken[0] - rest = [{k: v for k, v in row.items() if k != only} for row in rows] - else: - gone = set(taken) - rest = [{k: v for k, v in row.items() if k not in gone} for row in rows] + return taken, matrices, _without_keys(rows, taken) + + +def split_float_vectors(java_db, type_name, rows, declared=None, undeclared=False): + """Take the float vectors of `rows` out of the JSON text of a bulk load (the JVM wheel). + + A list of floats costs Python about 450 ns a number to write as text and the + engine about as much again to parse it, so a vector load spent most of its + time on the text. The vectors of the columns that take_vector_columns picks + cross as one float32 matrix per property instead (``FloatVectors.Matrix``, + one buffer copy). Every other property, and a column that is missing in a + row, ragged, not numeric, or null in a row, stays in the JSON. + + Returns ``(rest_rows, names, matrices)``: `rest_rows` are the rows without + those properties, `names` a ``String[]`` and `matrices` an ``Object[]`` of + ``FloatVectors.Matrix`` for the bridge; or None when no property qualifies + (the caller sends `rows` as they are). + """ + if _np is None or not rows: + return None + vectors = _float_vectors_class() + if vectors is None: + return None + if declared is None: + declared = declared_float_vectors(java_db, type_name) + taken = take_vector_columns(rows, declared, undeclared) + if taken is None: + return None + names, matrices, rest = taken + float_array = jpype.JArray(jpype.JFloat) return ( rest, - jpype.JArray(jpype.JString)(taken), - jpype.JArray(jpype.JObject)(matrices), + jpype.JArray(jpype.JString)(names), + jpype.JArray(jpype.JObject)( + [ + vectors.Matrix(float_array(matrix.reshape(-1)), int(matrix.shape[1])) + for matrix in matrices + ] + ), ) diff --git a/tests/conftest.py b/tests/conftest.py index ec3147c7..27d52e9e 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") + # the crossings of GraphBatch.create_vertices with float32 columns (lib.call, the codec's per-row hooks), observed on the native library + collect_ignore.append("test_native_vertex_vector_crossings.py") if not os.path.isfile( os.path.join(os.path.dirname(os.path.dirname(__file__)), "docs", "index.md") ): @@ -119,6 +121,8 @@ "test_query_vector_params_jvm.py": "the java.util objects the JVM build's parameter conversion builds", "test_insert_many_gc_pause.py": "the cyclic collector pause is the JVM build's (the native insert_many allocates nothing the collector tracks)", "test_graph_batch_tagged_shapes.py": "the JVM wheel's graph batch bulk path (bridge VertexBatcher and EdgeBatcher, the property-matrix fallback)", + "test_graph_batch_single_pass.py": "the JVM wheel's graph batch bulk path (GraphBatch._create_vertices_json_bulk, the property-matrix fallback)", + "test_graph_batch_vector_crossings_jvm.py": "the JPype crossings of the JVM wheel's create_vertices vector path (the bridge classes, the property-matrix fallback)", "test_async_proxy_cache.py": "the JPype proxy classes of the JVM build's async executor", "test_read_crossings.py": "crossing counts of the JVM build's read path (JPype proxies of the Java objects and the bridge classes)", "test_reader_crossings.py": "crossing counts of the JVM build's record and row readers (JPype proxies of the Java objects and the bridge classes)", diff --git a/tests/test_bulk_json_paths.py b/tests/test_bulk_json_paths.py index 8bb7bb07..1c801d00 100644 --- a/tests/test_bulk_json_paths.py +++ b/tests/test_bulk_json_paths.py @@ -104,7 +104,7 @@ def test_a_lone_surrogate_raises_instead_of_becoming_a_question_mark( _write(db, entry, "a\ud800b") -def test_ordinary_rows_stay_on_the_bulk_path(temp_db_path): +def test_ordinary_rows_stay_on_the_bulk_path(): rows = [ {"id": i, "name": f"n{i}", "score": i * 0.5, "flag": i % 2 == 0, "note": None} for i in range(5) @@ -113,14 +113,7 @@ def test_ordinary_rows_stay_on_the_bulk_path(temp_db_path): assert all( type_conversion.json_bulk_scalar_ok(v) for row in rows for v in row.values() ) - with arcadedb.create_database(temp_db_path) as db: - _create_types(db) - with db.graph_batch(parallel_flush=False) as batch: - assert len(batch._create_vertices_json_bulk("V", rows)) == len(rows) - assert batch._create_vertices_json_bulk("V", rows + [{"id": 2**70}]) is None - assert ( - batch._create_vertices_json_bulk("V", rows + [{"x": math.nan}]) is None - ) + # the graph batch's side of this (the JVM wheel's _create_vertices_json_bulk) is in test_graph_batch_single_pass.py @pytest.mark.parametrize( diff --git a/tests/test_bulk_single_pass_validation.py b/tests/test_bulk_single_pass_validation.py new file mode 100644 index 00000000..5d9afe72 --- /dev/null +++ b/tests/test_bulk_single_pass_validation.py @@ -0,0 +1,419 @@ +"""insert_many and create_vertices validate their rows in the one json.dumps pass, not in a Python loop over the rows before it. + +`json_bulk_dumps` and `json_bulk_dumps_tagged` used to walk every value of every row before json.dumps walked it again. Now json.dumps +writes the rows once and the check comes from what that pass shows (the text and the types of the rows), and a walk runs only when it +may find something. The contract is unchanged: the same text for rows that carry, and for a row that does not, the same exception type +and message as the old loop gave for the first bad row in row order. The old loop is copied here as the private reference, so these +tests fail on the old code (they spy on `_check_rows`, which only the new code has) and compare the new code with that reference. +The old messages name no row index, so the first bad row of a large batch is told from a later one by the value in its message. +""" + +import datetime as dt +import decimal +import enum +import json +import math +import random + +import numpy as np +import pytest + +import arcadedb +from arcadedb import type_conversion as tc + +INT64_MAX = 2**63 - 1 +INT64_MIN = -(2**63) + + +def _old_dumps(rows, tagged): + """The json_bulk_dumps and json_bulk_dumps_tagged of 26.11.1: the validation loop, then json.dumps.""" + for row in rows: + for key in row: + if type(key) is not str: + raise ValueError("a non-str key cannot go through the JSON bulk path") + for value in row.values(): + kind = type(value) + if kind is str: + if not value.isascii() and not tc._str_survives_json(value): + raise ValueError( + "a string with a lone surrogate cannot go through the JSON bulk path" + ) + elif kind is int: + if not INT64_MIN <= value <= INT64_MAX: + raise ValueError( + f"an integer beyond 64 bits cannot go through the JSON bulk path: {value}" + ) + elif kind is float or kind is bool or value is None: + continue + elif tagged and kind in tc._TAGGED_KINDS: + continue + elif tagged and kind is dict and tc._TAG in value: + raise ValueError( + "a map keyed by the tag marker cannot go through the JSON bulk path" + ) + elif tagged and ( + isinstance( + value, (dt.datetime, dt.date, decimal.Decimal, bytes, bytearray) + ) + or isinstance(value, np.generic) + or tc._is_pandas_null(kind) + ): + continue + else: + tc._json_bulk_check(value) + if tagged: + return tc._dumps_checked(rows, default=tc.bulk_tag_default) + return tc._dumps_checked(rows) + + +def _outcome(call, rows): + try: + return ("text", call(rows)) + except ( + BaseException + ) as e: # noqa: BLE001 - the type and the message are what is compared + return (type(e).__name__, str(e)) + + +class Colour(enum.IntEnum): + RED = 5 + + +class Name(str): + pass + + +def _scalar_rows(): + return [ + {"id": 1, "name": "alpha", "score": 0.5, "flag": True, "note": None}, + {"id": 2, "name": "", "score": -1e-30, "flag": False, "note": "x"}, + ] + + +CASES = { + "scalars": _scalar_rows(), + "empty": [], + "empty_rows": [{}, {}], + "bounds_fit": [{"a": INT64_MAX, "b": INT64_MIN, "c": 0}], + "above_max": [{"a": INT64_MAX + 1}], + "below_min": [{"a": INT64_MIN - 1}], + "big_negative": [{"a": -(10**30)}], + "huge_digits": [{"a": 10**5000}], + "non_bmp": [{"a": "emoji \U0001f600 pair", "b": "\U00010000"}], + "cjk_and_hangul": [{"a": "中文한글퟿"}], + "nul_and_controls": [{"a": "a\x00b\x1f", "b": "
"}], + "lone_high": [{"a": "a\ud800b"}], + "lone_low": [{"a": "\udc00"}], + "two_lone_halves": [{"a": "𐀀"}], + "literal_backslash_u": [{"a": "\\ud800", "b": "\\u0041"}], + "digits_in_text": [{"a": "9999999999999999999", "b": "1234567890123456789012"}], + "nan": [{"a": math.nan}], + "inf": [{"a": -math.inf}], + "subnormal": [{"a": 5e-324}], + "small_floats": [{"a": 0.00012345678901234568, "b": 1.2345678901234567e19}], + "int_key": [{1: "a"}], + "bool_key": [{True: "a"}], + "none_key": [{None: "a"}], + "list_ok": [{"a": [1, "x", None, [2.5]]}], + "list_big": [{"a": [1, 2**70]}], + "list_nan": [{"a": [1.0, math.nan]}], + "list_lone": [{"a": ["\ud800"]}], + "list_datetime": [{"a": [dt.datetime(2026, 1, 2)]}], + "dict_ok": [{"a": {"b": [1], "c": "d"}}], + "dict_int_key": [{"a": {1: 2}}], + "dict_tag_key": [{"a": {tc._TAG: 1}}], + "tuple_value": [{"a": (1, 2)}], + "datetime": [{"a": dt.datetime(2026, 1, 2, 3, 4, 5, 6)}], + "aware_datetime": [ + {"a": dt.datetime(2026, 1, 2, tzinfo=dt.timezone(dt.timedelta(hours=9)))} + ], + "date": [{"a": dt.date(2026, 1, 2)}], + "decimal": [{"a": decimal.Decimal("1.50")}], + "decimal_nan": [{"a": decimal.Decimal("NaN")}], + "bytes": [{"a": b"\x00\xffab", "b": bytearray(b"x")}], + "numpy_scalars": [ + { + "a": np.int64(5), + "b": np.float64(1.5), + "c": np.bool_(True), + "d": np.float32(2.5), + } + ], + "numpy_big": [{"a": np.uint64(2**64 - 1)}], + "numpy_nan": [{"a": np.float64(math.nan)}], + "numpy_datetime": [{"a": np.datetime64("2026-01-02T03:04:05")}], + "numpy_str": [{"a": np.str_("q")}], + "int_enum": [{"a": Colour.RED}], + "str_subclass": [{"a": Name("s")}], + "str_subclass_lone": [{"a": Name("\ud800")}], + "set_value": [{"a": {1, 2}}], + "object_value": [{"a": object()}], + "none_row": [None], + "list_row": [[1, 2]], + "string_row": ["ab"], + "mixed_good_then_bad": [{"a": 1}, {"a": 2**64}, {"a": math.nan}], +} + + +@pytest.mark.parametrize("tagged", [False, True], ids=["untagged", "tagged"]) +@pytest.mark.parametrize("name", sorted(CASES)) +def test_same_text_or_same_error_as_the_validation_loop(name, tagged): + rows = CASES[name] + new = tc.json_bulk_dumps_tagged if tagged else tc.json_bulk_dumps + assert _outcome(new, rows) == _outcome(lambda r: _old_dumps(r, tagged), rows) + + +@pytest.mark.parametrize("tagged", [False, True], ids=["untagged", "tagged"]) +def test_random_rows_agree_with_the_validation_loop(tagged): + atoms = [ + None, True, 0, -1, INT64_MAX, INT64_MIN, INT64_MAX + 1, 10**30, 1.5, 0.00012345678901234568, + 5e-324, math.nan, math.inf, "", "abc", "a\x00b", "é中", "\U0001f600", "a\ud800b", "\udc00", + "𐀀", "9999999999999999999", [1, 2], [math.nan], {"a": 1}, {1: 2}, (1, 2), + dt.datetime(2026, 1, 2), dt.date(2026, 1, 2), decimal.Decimal("1.5"), b"ab", np.int64(7), + np.float64(math.nan), Colour.RED, Name("s"), {1, 2}, object(), + ] # fmt: skip + rnd = random.Random(318) # nosec B311 - a seeded corpus, not security + new = tc.json_bulk_dumps_tagged if tagged else tc.json_bulk_dumps + for _ in range(3000): + rows = [] + for _ in range(rnd.choice([1, 2, 3, 8])): + keys = rnd.sample(["a", "b", "c", "d"], rnd.randint(0, 4)) + if rnd.random() < 0.03: + keys.append(1) + rows.append({k: rnd.choice(atoms) for k in keys}) + assert _outcome(new, rows) == _outcome(lambda r: _old_dumps(r, tagged), rows) + + +def _batch(n, bad): + """n scalar rows of the shape of a TPC-H part table; `bad` maps a row index to its replacement row.""" + rows = [ + { + "p_partkey": i, + "p_name": f"part {i}", + "p_size": i % 50, + "p_retailprice": 900.0 + (i % 1000) / 100, + "p_comment": None if i % 7 == 0 else "carefully final deposits", + } + for i in range(n) + ] + for index, row in bad.items(): + rows[index] = row + return rows + + +# kind -> (first bad row, a later bad row of another value), so the message tells which row raised +BAD_KINDS = { + "int_beyond_64_bits": ({"p_size": 2**63 + 11}, {"p_size": 2**70 + 22}), + "negative_beyond_64_bits": ({"p_size": INT64_MIN - 33}, {"p_size": 2**64 + 44}), + "lone_surrogate": ({"p_name": "a\ud800"}, {"p_size": 2**65}), + "non_str_key": ({7: "x"}, {"p_size": 2**65}), + "nested_int": ({"p_name": [1, 2**66]}, {"p_size": 2**65 + 55}), + "nan": ({"p_retailprice": math.nan}, {"p_size": 2**65}), + "subnormal": ({"p_retailprice": 5e-324}, {"p_size": 2**65}), +} + + +@pytest.mark.parametrize("tagged", [False, True], ids=["untagged", "tagged"]) +@pytest.mark.parametrize("kind", sorted(BAD_KINDS)) +def test_the_first_bad_row_of_a_large_batch_gives_the_old_error(kind, tagged): + first, later = BAD_KINDS[kind] + rows = _batch(20_000, {12_345: first, 19_000: later}) + new = tc.json_bulk_dumps_tagged if tagged else tc.json_bulk_dumps + expected = _outcome(lambda r: _old_dumps(r, tagged), rows) + assert expected[0] in ("ValueError", "TypeError"), expected + assert _outcome(new, rows) == expected + # the later row alone raises a message of its own. A NaN or a subnormal is refused by json.dumps, after the old loop had walked + # every row, so the old error there is the later row's; for every other kind the first bad row's error comes first + only_later = _batch(20_000, {19_000: later}) + assert _outcome(new, only_later) == _outcome( + lambda r: _old_dumps(r, tagged), only_later + ) + assert (_outcome(new, only_later) == expected) == (kind in ("nan", "subnormal")) + + +@pytest.mark.parametrize("tagged", [False, True], ids=["untagged", "tagged"]) +def test_a_clean_large_batch_is_the_plain_json_text(tagged): + rows = _batch(20_000, {}) + new = tc.json_bulk_dumps_tagged if tagged else tc.json_bulk_dumps + assert new(rows) == json.dumps(rows) + + +def test_scalar_rows_are_not_walked_when_the_text_carries_them(monkeypatch): + # the insert_many shape of the paper's l1tpc Part load and the create_vertices shape of l2, e2 and l3d: dict rows of str, int, float, + # bool and None. Numbers that look alike but fit: 19-digit ints below 2**63, floats whose text has a long run of digits after the + # point, text with digits, Hangul and CJK (BMP characters that json.dumps writes as \uXXXX, some of them as \ud...). + walks = [] + real = tc._check_rows + monkeypatch.setattr( + tc, + "_check_rows", + lambda rows, tagged: walks.append(tagged) or real(rows, tagged), + ) + rows = _batch(2_000, {}) + rows += [ + { + "ts": 1_700_000_000_123_456_789, + "hi": INT64_MAX, + "lo": INT64_MIN + 1, + "f": 0.00012345678901234568, + }, + { + "f": 1.2345678901234567e19, + "s": "id-123456789012345678", + "han": "한글퟿中文", + "nul": "a\x00b", + }, + {"e": "", "n": None, "b": True, "neg": -123456789012345678}, + ] + assert tc.json_bulk_dumps(rows) == json.dumps(rows) + assert tc.json_bulk_dumps_tagged(rows) == json.dumps(rows) + assert walks == [] + + +def test_tagged_kinds_are_not_walked_either(monkeypatch): + walks = [] + real = tc._check_rows + monkeypatch.setattr( + tc, + "_check_rows", + lambda rows, tagged: walks.append(tagged) or real(rows, tagged), + ) + rows = [ + {"d": dt.datetime(2026, 1, 2, 3, 4, 5), "day": dt.date(2026, 1, 2), "m": decimal.Decimal("1.50"), "b": b"ab", "n": np.int64(5), "f": np.float32(2.5)} + for _ in range(100) + ] # fmt: skip + assert tc.json_bulk_dumps_tagged(rows) + assert walks == [] + + +@pytest.mark.parametrize( + "row", + [ + {"a": 2**63}, + {"a": 10**20}, + {"a": INT64_MIN - 1}, + {"a": "a\ud800"}, + {"a": "\U0001f600"}, + {7: "x"}, + {"a": [1]}, + {"a": 5e-324}, + {"a": Colour.RED}, + ], +) +def test_a_value_that_may_be_bad_runs_the_walk(monkeypatch, row): + walks = [] + real = tc._check_rows + monkeypatch.setattr( + tc, + "_check_rows", + lambda rows, tagged: walks.append(tagged) or real(rows, tagged), + ) + _outcome(tc.json_bulk_dumps, [{"id": 1}, row]) + assert walks, "the walk must decide a row the text cannot vouch for" + + +def test_int64_overflow_pattern_matches_exactly_the_tokens_from_2_to_the_63(): + rnd = random.Random(63) # nosec B311 - a seeded corpus, not security + values = [ + 0, + 1, + 10**18, + 2**63 - 2, + 2**63 - 1, + 2**63, + 2**63 + 1, + 10**19 - 1, + 10**19, + 10**20, + 9 * 10**18, + 9_999_999_999_999_999_999, + ] + values += [rnd.randint(10**18, 2**65) for _ in range(5000)] # fmt: skip + for value in values: + for text in (str(value), str(-value), f"[{value},1]", f'{{"a":{value}}}'): + hit = tc._INT64_OVERFLOW.search(text) is not None + assert hit == (value >= 2**63), text + # digits after the point of a float and inside an exponent are no integer token + assert tc._INT64_OVERFLOW.search("[0.00012345678901234568]") is None + assert tc._INT64_OVERFLOW.search("[1.2345678901234567e+19]") is None + + +# --------------------------------------------------------------------------- +# the entry points, on whichever wheel is installed + + +def _types(db): + db.command("sql", "CREATE DOCUMENT TYPE Part") + db.command("sql", "CREATE VERTEX TYPE PartV") + + +def _stored(db, type_name): + query = f"SELECT FROM {type_name}" # nosec B608 - test-controlled + rows = db.query("sql", query).to_list() + return sorted( + ( + r["p_partkey"], + r["p_name"], + r["p_size"], + r["p_retailprice"], + r.get("p_comment"), + ) + for r in rows + ) + + +def _expected(rows): + return sorted( + ( + r["p_partkey"], + r["p_name"], + r["p_size"], + r["p_retailprice"], + r.get("p_comment"), + ) + for r in rows + ) + + +def _rows_with_text(): + rows = _batch(25_000, {}) + rows[10]["p_name"] = "emoji \U0001f600 \U00010000" + rows[11]["p_name"] = "" + rows[12]["p_name"] = "한글 中文 ퟿" + rows[13]["p_size"] = INT64_MAX + rows[14]["p_size"] = INT64_MIN + return rows + + +def test_insert_many_stores_a_scalar_batch_as_before(temp_db_path): + rows = _rows_with_text() + with arcadedb.create_database(temp_db_path) as db: + _types(db) + assert db.insert_many("Part", rows) == len(rows) + assert _stored(db, "Part") == _expected(rows) + + +def test_create_vertices_stores_a_scalar_batch_as_before(temp_db_path): + rows = _rows_with_text() + with arcadedb.create_database(temp_db_path) as db: + _types(db) + with db.graph_batch(parallel_flush=False) as batch: + assert len(batch.create_vertices("PartV", rows)) == len(rows) + assert _stored(db, "PartV") == _expected(rows) + + +@pytest.mark.parametrize("entry", ["insert_many", "create_vertices"]) +@pytest.mark.parametrize("bad", [2**63 + 11, "a\ud800"], ids=["int", "surrogate"]) +def test_a_bad_row_inside_a_large_batch_is_refused(temp_db_path, entry, bad): + column = "p_size" if isinstance(bad, int) else "p_name" + rows = _batch(25_000, {12_345: {column: bad}, 19_000: {"p_size": 2**70 + 22}}) + with arcadedb.create_database(temp_db_path) as db: + _types(db) + with pytest.raises( + Exception + ): # noqa: B017 - the entry points raise their own types + if entry == "insert_many": + db.insert_many("Part", rows) + else: + with db.graph_batch(parallel_flush=False) as batch: + batch.create_vertices("PartV", rows) diff --git a/tests/test_graph_batch_single_pass.py b/tests/test_graph_batch_single_pass.py new file mode 100644 index 00000000..2415c012 --- /dev/null +++ b/tests/test_graph_batch_single_pass.py @@ -0,0 +1,84 @@ +"""GraphBatch.create_vertices (the JVM wheel) decides its rows by their types and writes each chunk once. + +Before, `_create_vertices_json_bulk` called `json_bulk_scalar_ok` on every value of every row before the first chunk was written, and +the chunk was then walked again by json.dumps. Now `type_conversion.classify_vertex_rows` takes a census of the types in C, and the +value checks come from the one json.dumps pass of each chunk (`json_bulk_dumps`, as `json_bulk_dumps_tagged` already did for the tagged +text). What stays: the whole input goes the property-matrix way for a property name that is not a str or a type the text does not carry +(`_create_vertices_json_bulk` returns None); what changes: a scalar the text would change (a subnormal, a NaN, an integer beyond 64 bits, +a lone surrogate) sends its own chunk, not the whole input, through the matrix path, as it already did for a tagged chunk. +""" + +import pytest + +import arcadedb +from arcadedb import graph_batch as gbm +from arcadedb import type_conversion as tc +from arcadedb.graph_batch import GraphBatch + + +def _rows(n): + return [ + {"id": i, "name": f"n{i}", "score": i * 0.5, "flag": i % 2 == 0, "note": None} + for i in range(n) + ] + + +def test_scalar_rows_are_not_judged_value_by_value(temp_db_path, monkeypatch): + def refuse(value): + raise AssertionError("a per-value gate ran over a scalar table") + + monkeypatch.setattr(tc, "json_bulk_scalar_ok", refuse) + monkeypatch.setattr(gbm, "json_bulk_scalar_ok", refuse) + monkeypatch.setattr(tc, "json_bulk_shape_ok", refuse) + monkeypatch.setattr(gbm, "json_bulk_shape_ok", refuse) + rows = _rows(2_000) + with arcadedb.create_database(temp_db_path) as db: + db.command("sql", "CREATE VERTEX TYPE V") + with db.graph_batch(parallel_flush=False) as batch: + assert len(batch._create_vertices_json_bulk("V", rows)) == len(rows) + count = db.query("sql", "SELECT count(*) AS c FROM V").first().get("c") + assert count == len(rows) + + +def test_the_whole_input_still_goes_the_matrix_way_for_a_name_or_type_the_text_does_not_carry( + temp_db_path, +): + rows = _rows(5) + with arcadedb.create_database(temp_db_path) as db: + db.command("sql", "CREATE VERTEX TYPE V") + with db.graph_batch(parallel_flush=False) as batch: + assert len(batch._create_vertices_json_bulk("V", rows)) == len(rows) + assert batch._create_vertices_json_bulk("V", rows + [{1: "a"}]) is None + assert batch._create_vertices_json_bulk("V", rows + [{"k": {1, 2}}]) is None + + +@pytest.mark.parametrize("value", [5e-324, float("nan"), 2**70, "a\ud800b"]) +def test_a_scalar_the_text_changes_sends_its_chunk_through_the_matrix_path( + temp_db_path, monkeypatch, value +): + monkeypatch.setattr(GraphBatch, "_BULK_CHUNK", 100) + calls = [] + real = GraphBatch._to_java_property_matrix + + def spy(self, chunk): + calls.append(len(chunk)) + return real(self, chunk) + + monkeypatch.setattr(GraphBatch, "_to_java_property_matrix", spy) + rows = [{"id": i, "x": 1.5} for i in range(250)] + rows[150] = {"id": 150, "x": value} + with arcadedb.create_database(temp_db_path) as db: + db.command("sql", "CREATE VERTEX TYPE V") + with db.graph_batch(parallel_flush=False) as batch: + try: + rids = batch.create_vertices("V", rows) + except ( + Exception + ): # noqa: BLE001 - the matrix path refuses a value that no entry point stores + rids = None + assert calls == [100], "only the chunk of row 150 takes the matrix path" + if value == 5e-324: + # stored as the double every other entry point stores; the text of a subnormal was read as a DECIMAL + assert len(rids) == 250 + stored = db.query("sql", "SELECT x FROM V WHERE id = 150").first().get("x") + assert stored == 5e-324 and type(stored) is float diff --git a/tests/test_graph_batch_tagged_shapes.py b/tests/test_graph_batch_tagged_shapes.py index 2da99c18..560b2d98 100644 --- a/tests/test_graph_batch_tagged_shapes.py +++ b/tests/test_graph_batch_tagged_shapes.py @@ -438,6 +438,8 @@ def shape(*a, **k): monkeypatch.setattr(gbm, "json_bulk_dumps_tagged", dumps) monkeypatch.setattr(gbm, "json_bulk_shape_ok", shape) + # create_vertices decides the shapes in type_conversion.classify_vertex_rows, which calls the gate there + monkeypatch.setattr(type_conversion, "json_bulk_shape_ok", shape) return calls diff --git a/tests/test_graph_batch_vector_crossings_jvm.py b/tests/test_graph_batch_vector_crossings_jvm.py new file mode 100644 index 00000000..a16d36c3 --- /dev/null +++ b/tests/test_graph_batch_vector_crossings_jvm.py @@ -0,0 +1,241 @@ +"""GraphBatch.create_vertices (the JVM wheel): a Java float[] or float32 array crosses JPype as one matrix per property and chunk (#337). + +Before, a Java ``float[]`` (what ``to_java_float_array`` returns) was no vector candidate: ``classify_vertex_rows`` gave up on the first +row, and the whole load went through ``_to_java_property_matrix``, which converts every value of every row with ``convert_python_to_java``, +builds one ``Object[]`` per row and one ``Object[][]`` for the call (an unwrapped Python loop that was 59 to 61 percent of the wall of the +e2 and l3d loads in the untimed ingest profile of 26.11.1). A float32 numpy array on a property that is not declared ``ARRAY_OF_FLOATS`` +took the same way. Now ``take_vector_columns`` stacks the column once (numpy reads a Java array through the buffer protocol), one +``float[]`` copy and one ``FloatVectors.Matrix`` cross per property, and one ``VertexBatcher.createVerticesJson`` call per chunk. + +The JPype crossings are counted by wrapping the bridge classes and the Python-side conversion; the stored values are in +test_vertex_vector_loads.py. The JPype bridge does not exist on the native wheel (its counts are in test_native_vertex_vector_crossings.py), +so conftest.py does not collect this file there. +""" + +import jpype +import numpy as np +import pytest + +import arcadedb +from arcadedb import graph_batch as gbm +from arcadedb import type_conversion as tc +from arcadedb.graph_batch import GraphBatch + +ROWS, CHUNK, DIM = 250, 100, 16 # three chunks: 100, 100, 50 + + +class _Counting: + """Forwards every call of a Java class (or of the bridge's FloatVectors) and records its name.""" + + def __init__(self, target, log): + self._target, self._log = target, log + + def __getattr__(self, name): + attribute = getattr(self._target, name) + if not callable(attribute): + return attribute + + def call(*args): + self._log.append(name) + return attribute(*args) + + return call + + +@pytest.fixture +def crossings(monkeypatch): + """What the load does: property-matrix fallbacks, per-value conversions, bridge calls, Matrix constructions, per-row inserts.""" + log = {"fallback": [], "convert": [], "bridge": [], "matrix": [], "per_row": []} + + real_matrix_path = GraphBatch._to_java_property_matrix + + def fallback(self, rows): + log["fallback"].append(len(rows)) + return real_matrix_path(self, rows) + + monkeypatch.setattr(GraphBatch, "_to_java_property_matrix", fallback) + + real_convert = tc.convert_python_to_java + + def convert(value): + log["convert"].append(type(value).__name__) + return real_convert(value) + + monkeypatch.setattr(tc, "convert_python_to_java", convert) + monkeypatch.setattr(gbm, "convert_python_to_java", convert) + + real_jclass = jpype.JClass + + def jclass(name, *args, **kwargs): + cls = real_jclass(name, *args, **kwargs) + if name == "com.arcadedb.python.VertexBatcher": + return _Counting(cls, log["bridge"]) + return cls + + monkeypatch.setattr(jpype, "JClass", jclass) + + real_vectors = tc._float_vectors_class + + def vectors_class(): + vectors = real_vectors() + return None if vectors is None else _Counting(vectors, log["matrix"]) + + monkeypatch.setattr(tc, "_float_vectors_class", vectors_class) + + from arcadedb import core + + real_per_row = core.Database._insert_many_per_row + + def per_row(self, type_name, rows): + log["per_row"].append(len(rows)) + return real_per_row(self, type_name, rows) + + monkeypatch.setattr(core.Database, "_insert_many_per_row", per_row) + monkeypatch.setattr(GraphBatch, "_BULK_CHUNK", CHUNK) + return log + + +@pytest.fixture +def db(temp_db): + for name in ("Declared", "Free", "Doc"): + kind = "DOCUMENT" if name == "Doc" else "VERTEX" + temp_db.command("sql", f"CREATE {kind} TYPE {name}") + temp_db.command("sql", f"CREATE PROPERTY {name}.id LONG") + temp_db.command("sql", "CREATE PROPERTY Declared.embedding ARRAY_OF_FLOATS") + temp_db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + return temp_db + + +def _x(n=ROWS, dim=DIM): + return np.random.default_rng(7).standard_normal((n, dim)).astype(np.float32) + + +def _floats(x): + return [jpype.JArray(jpype.JFloat)(row) for row in x] + + +@pytest.mark.parametrize("type_name", ["Declared", "Free"]) +@pytest.mark.parametrize("shape", ["java_float_array", "float32_ndarray"]) +def test_a_float32_vector_crosses_as_one_matrix_per_chunk( + db, crossings, type_name, shape +): + x = _x() + vectors = _floats(x) if shape == "java_float_array" else list(x) + rows = [{"id": i, "embedding": vectors[i]} for i in range(ROWS)] + with db.graph_batch(parallel_flush=False) as gb: + assert len(gb.create_vertices(type_name, rows)) == ROWS + assert ( + crossings["fallback"] == [] + ), "the property-matrix path ran (a JPype conversion per value and per row)" + assert crossings["convert"] == [], "a value was converted one at a time" + assert ( + crossings["bridge"] == ["createVerticesJson"] * 3 + ) # one crossing per chunk of 100 + # one float32 matrix per chunk: `declared` is asked once per call, `Matrix` once per chunk + assert crossings["matrix"].count("Matrix") == 3 + sql = f"SELECT id, embedding FROM {type_name}" # nosec B608 - test-owned name + stored = {r["id"]: r["embedding"] for r in db.query("sql", sql).to_list()} + for i in range(ROWS): + assert np.array(stored[i], dtype=np.float32).tobytes() == x[i].tobytes() + + +def test_two_vector_properties_cross_as_two_matrices_per_chunk(db, crossings): + x, y = _x(), _x(dim=8) + rows = [ + {"id": i, "embedding": jpype.JArray(jpype.JFloat)(x[i]), "second": y[i]} + for i in range(ROWS) + ] + with db.graph_batch(parallel_flush=False) as gb: + gb.create_vertices("Declared", rows) + assert crossings["fallback"] == [] and crossings["convert"] == [] + assert crossings["bridge"] == ["createVerticesJson"] * 3 + assert crossings["matrix"].count("Matrix") == 6 + + +def test_a_ragged_column_still_takes_the_property_matrix_path_for_its_chunk( + db, crossings +): + x = _x() + vectors = _floats(x) + rows = [{"id": i, "embedding": vectors[i]} for i in range(ROWS)] + rows[120] = {"id": 120, "embedding": jpype.JArray(jpype.JFloat)(x[120][:5])} + with db.graph_batch(parallel_flush=False) as gb: + assert len(gb.create_vertices("Declared", rows)) == ROWS + assert crossings["fallback"] == [100], "only the chunk of row 120" + assert crossings["bridge"] == ["createVerticesJson"] * 2 + stored = { + r["id"]: r["embedding"] + for r in db.query("sql", "SELECT id, embedding FROM Declared").to_list() + } + assert np.array(stored[120], dtype=np.float32).tobytes() == x[120][:5].tobytes() + assert np.array(stored[0], dtype=np.float32).tobytes() == x[0].tobytes() + + +def test_a_java_double_array_is_a_vector_only_on_a_declared_property(db, crossings): + x64 = np.random.default_rng(8).random((ROWS, DIM)) * 1000.0 + 1e-3 + rows = [ + {"id": i, "embedding": jpype.JArray(jpype.JDouble)(x64[i])} for i in range(ROWS) + ] + with db.graph_batch(parallel_flush=False) as gb: + gb.create_vertices("Declared", rows) + assert ( + crossings["fallback"] == [] + and crossings["bridge"] == ["createVerticesJson"] * 3 + ) + declared = { + r["id"]: r["embedding"] + for r in db.query("sql", "SELECT id, embedding FROM Declared").to_list() + } + for i in ( + 0, + 99, + 100, + ROWS - 1, + ): # the engine stores floats for the declared property + want = x64[i].astype(np.float32).tobytes() + assert np.array(declared[i], dtype=np.float32).tobytes() == want + del crossings["fallback"][:], crossings["bridge"][:] + with db.graph_batch(parallel_flush=False) as gb: + gb.create_vertices("Free", rows) + assert crossings["fallback"] == [ + ROWS + ], "undeclared: the double[] keeps the matrix path" + free = { + r["id"]: r["embedding"] + for r in db.query("sql", "SELECT id, embedding FROM Free").to_list() + } + assert np.array(free[5], dtype=np.float64).tolist() == x64[5].tolist() + + +def test_insert_many_sends_a_java_float_array_column_through_the_matrix(db, crossings): + x = _x(50) + arrays = _floats(x) + rows = [{"id": i, "embedding": arrays[i]} for i in range(50)] + assert db.insert_many("Doc", rows) == 50 + assert crossings["per_row"] == [], "the column went row by row" + assert crossings["matrix"].count("Matrix") == 1 + stored = { + r["id"]: r["embedding"] + for r in db.query("sql", "SELECT id, embedding FROM Doc").to_list() + } + for i in range(50): + assert np.array(stored[i], dtype=np.float32).tobytes() == x[i].tobytes() + + +def test_take_vector_columns_reads_java_arrays_through_one_stack(db): + float_array, double_array = jpype.JArray(jpype.JFloat), jpype.JArray(jpype.JDouble) + x = _x(6, 5) + floats = [{"v": float_array(row), "k": i} for i, row in enumerate(x)] + names, matrices, rest = tc.take_vector_columns(floats, (), undeclared=True) + assert names == ["v"] and matrices[0].dtype == np.float32 + assert matrices[0].tobytes() == x.tobytes() and rest == [{"k": i} for i in range(6)] + doubles = [{"v": double_array(row.astype(np.float64))} for row in x] + assert tc.take_vector_columns(doubles, (), undeclared=True) is None + assert tc.take_vector_columns(doubles, ["v"])[1][0].tobytes() == x.tobytes() + ints = [{"v": jpype.JArray(jpype.JInt)([1, 2, 3])} for _ in range(3)] + assert tc.take_vector_columns(ints, ["v"], undeclared=True) is None + assert not tc.is_vector_candidate(ints[0]["v"]) + assert tc.is_vector_candidate(floats[0]["v"]) and tc.is_vector_candidate( + doubles[0]["v"] + ) + assert isinstance(arcadedb.to_java_float_array(x[0]), float_array) diff --git a/tests/test_native_graph_batch_binary_vectors.py b/tests/test_native_graph_batch_binary_vectors.py index 429adf49..c5543109 100644 --- a/tests/test_native_graph_batch_binary_vectors.py +++ b/tests/test_native_graph_batch_binary_vectors.py @@ -1,7 +1,8 @@ """GraphBatch float32 vectors cross the native library in a binary side buffer, not as number text. ``create_vertices`` rows and ``new_edges`` props that hold a float32 ``ndarray`` go through ``arcadedb_n_gb_bin`` (the float32 buffer is -read in Java through the ``["B", offset, count]`` tag, as ``arcadedb_n_exec_bin`` reads query parameters); everything else keeps the text +read in Java through the ``["B", offset, count]`` tag, as ``arcadedb_n_exec_bin`` reads query parameters; a ``create_vertices`` column of +float32 arrays of one length crosses as one matrix described by ``"vecs"`` instead, see test_native_vertex_vector_crossings.py); everything else keeps the text of ``arcadedb_n_gb``: a vector given as a list, a float64 array, a row with no vector. The stored values do not change: the same float32 as the ``"F"``-tagged text gave, bit for bit. A library without the entry (``lib.GB_BIN`` false) keeps the text path. @@ -95,9 +96,11 @@ def test_float32_vectors_take_the_binary_entry_and_are_stored_bit_equal(db, call assert calls.entries("create_vertices") == ["arcadedb_n_gb_bin"] _, _, text, count = next(c for c in calls.calls if c[1] == "create_vertices") assert count == N * dim - assert '"B"' in text and '"F"' not in text + assert ( + '"vecs"' in text and '"B"' not in text and '"F"' not in text + ) # one matrix, no tag per row (#337) assert str(float(x[2, 0])) not in text # the numbers are not in the text - assert len(text) < N * 80 # a vector is a tag, not dim numbers + assert len(text) < N * 80 # a row is its other properties, not dim numbers _assert_bit_equal(_stored(db, "Vx"), x) diff --git a/tests/test_native_vertex_vector_crossings.py b/tests/test_native_vertex_vector_crossings.py new file mode 100644 index 00000000..e6848451 --- /dev/null +++ b/tests/test_native_vertex_vector_crossings.py @@ -0,0 +1,256 @@ +"""Native GraphBatch.create_vertices: float32 vectors are split and encoded as columns, not row by row (#337). + +Before, every float32 array of a row was found by ``json.dumps`` through a Python hook (``codec._gb_default``, then ``_default``: about a +dozen ``isinstance`` tests and one array copy per row), the rows were walked by ``codec._prep`` first (one Python call per value), a tag +``{"\\u0001":["B",offset,count]}`` was written and parsed back in Java per row, and ``np.concatenate`` joined one small array per row +(the 29 to 31 percent prep share of the e2 and l3d loads in the untimed ingest profile of 26.11.1). Now ``take_vector_columns`` stacks each +float32 column once, ``codec.encode_vertex_args`` writes the other properties with one ``json.dumps`` (no hook when they are plain scalars) +and the matrices cross as the front of the side buffer, described by one ``"vecs"`` entry, which the library +(``NativeExtras.withVectors``) turns into one ``float[]`` per row. One ``arcadedb_n_gb_bin`` crossing per chunk, as before. + +Needs the native library, so conftest.py collects this file on the native backend only. The stored values are in test_vertex_vector_loads.py +(both wheels) and test_native_graph_batch_binary_vectors.py. +""" + +import ctypes +import json +from datetime import datetime + +import numpy as np +import pytest + +from arcadedb._native import codec, lib + +CHUNK, ROWS, DIM = 100, 250, 16 # three chunks: 100, 100, 50 + + +class _Spy: + """The create_vertices crossings (entry, JSON text, side buffer, float count) and the per-row Python calls of the codec.""" + + def __init__(self, monkeypatch): + self.crossings, self.hooks = [], {"gb_default": 0, "default": 0, "prep": 0} + real_call, real_gb_default = lib.call, codec._gb_default + real_default, real_prep = codec._default, codec._prep + + def call(name, *args): + if ( + name in ("arcadedb_n_gb", "arcadedb_n_gb_bin") + and args[1] == b"create_vertices" + ): + buf = args[3] if name == "arcadedb_n_gb_bin" else None + count = args[4] if name == "arcadedb_n_gb_bin" else 0 + self.crossings.append((name, args[2].decode(), buf, count)) + return real_call(name, *args) + + def counting(key, real): + def wrapper(*args, **kwargs): + self.hooks[key] += 1 + return real(*args, **kwargs) + + return wrapper + + monkeypatch.setattr(lib, "call", call) + monkeypatch.setattr( + codec, "_gb_default", counting("gb_default", real_gb_default) + ) + monkeypatch.setattr(codec, "_default", counting("default", real_default)) + monkeypatch.setattr(codec, "_prep", counting("prep", real_prep)) + + def entries(self): + return [c[0] for c in self.crossings] + + def around(self, call): + """``(result, hooks, crossings)`` of one ``call()``: the codec hook calls and the library crossings made inside it only. + + ``codec._prep`` is counted by name, so it also counts the calls the batch itself makes outside ``create_vertices``: creating + the batch (``_enc`` of its options, whose ``_prep`` recurses into the patched name once per option) and ``close()`` + (``encode_gb_args({})``). Those are a few calls and do not depend on the rows, so the counts here are taken around the + ``create_vertices`` call alone. + """ + before, seen = dict(self.hooks), len(self.crossings) + result = call() + used = {key: self.hooks[key] - before[key] for key in before} + return result, used, len(self.crossings) - seen + + +@pytest.fixture +def spy(monkeypatch): + return _Spy(monkeypatch) + + +@pytest.fixture +def db(temp_db): + for name in ("Declared", "Free"): + temp_db.command("sql", f"CREATE VERTEX TYPE {name}") + temp_db.command("sql", f"CREATE PROPERTY {name}.id LONG") + temp_db.command("sql", "CREATE PROPERTY Declared.embedding ARRAY_OF_FLOATS") + temp_db.command("sql", "CREATE PROPERTY Free.born DATETIME") + return temp_db + + +def _x(n=ROWS, dim=DIM, seed=7): + return np.random.default_rng(seed).standard_normal((n, dim)).astype(np.float32) + + +def _stored(db, type_name, prop="embedding"): + sql = f"SELECT id, {prop} AS v FROM {type_name}" # nosec B608 - test-owned names + return {r["id"]: r["v"] for r in db.query("sql", sql).to_list()} + + +def _assert_bit_equal(stored, x): + assert sorted(stored) == list(range(len(x))) + for i in range(len(x)): + assert np.array(stored[i], dtype=np.float32).tobytes() == x[i].tobytes() + + +@pytest.mark.parametrize("type_name", ["Declared", "Free"]) +def test_a_float32_column_is_one_matrix_with_no_hook_call_per_row(db, spy, type_name): + x = _x() + rows = [{"id": i, "embedding": x[i]} for i in range(ROWS)] + with db.graph_batch(commit_every=CHUNK) as gb: + rids, used, crossed = spy.around(lambda: gb.create_vertices(type_name, rows)) + assert len(rids) == ROWS + assert spy.entries() == ["arcadedb_n_gb_bin"] * 3 # one crossing per chunk of 100 + # no per-row hook: the per-row hooks of the old path are never called, and nothing is called more than once per crossing + assert used["gb_default"] == 0 and used["default"] == 0, used + assert used["prep"] <= crossed, used # the old path called it for every row + assert [c[3] for c in spy.crossings] == [CHUNK * DIM, CHUNK * DIM, 50 * DIM] + for _, text, buf, count in spy.crossings: + args = json.loads(text) + assert args["vecs"] == {"names": ["embedding"], "dims": [DIM], "offs": [0]} + assert '"B"' not in text and '"F"' not in text # no tag and no number text + assert all(set(row) == {"id"} for row in args["rows"]) + assert isinstance( + buf, ctypes.Array + ) # a view of the matrix, not a copy of its bytes + _assert_bit_equal(_stored(db, type_name), x) + + +def test_the_python_work_does_not_grow_with_the_row_count(db, spy): + # 250 rows in three chunks of 100, then 2,500 rows in three chunks of 1,000: ten times the rows, the same number of crossings, and + # the same bound on the codec hooks (the old path called the hooks once or more for every row, so it would show 250 and 2,500) + for rows_n, chunk in ((250, 100), (2500, 1000)): + x = _x(rows_n) + rows = [{"id": i, "embedding": x[i]} for i in range(rows_n)] + with db.graph_batch(commit_every=chunk) as gb: + rids, used, crossed = spy.around(lambda: gb.create_vertices("Free", rows)) + assert len(rids) == rows_n and crossed == 3 + assert used["gb_default"] == 0 and used["default"] == 0, used + assert used["prep"] <= crossed, used + + +def test_two_columns_and_a_date_time_keep_the_buffer_layout(db, spy): + x, y = _x(), _x(dim=8, seed=8) + when = datetime(2020, 1, 2, 3, 4, 5) + rows = [ + {"id": i, "embedding": x[i], "second": y[i], "born": when} for i in range(ROWS) + ] + with db.graph_batch(commit_every=CHUNK) as gb: + gb.create_vertices("Free", rows) + assert spy.entries() == ["arcadedb_n_gb_bin"] * 3 + first = json.loads(spy.crossings[0][1]) + assert first["vecs"] == { + "names": ["embedding", "second"], + "dims": [DIM, 8], + "offs": [0, CHUNK * DIM], + } + assert spy.crossings[0][3] == CHUNK * DIM + CHUNK * 8 + _assert_bit_equal(_stored(db, "Free"), x) + _assert_bit_equal(_stored(db, "Free", "second"), y) + sql = "SELECT born FROM Free WHERE id = 7" + assert db.query("sql", sql).first().get("born") == when + + +def test_a_float32_array_that_does_not_qualify_follows_the_matrices_in_the_buffer( + db, spy +): + x = _x(60) + rows = [{"id": i, "embedding": x[i], "short": x[i][: 2 + i % 2]} for i in range(60)] + with db.graph_batch(commit_every=CHUNK) as gb: + gb.create_vertices("Free", rows) + assert spy.entries() == ["arcadedb_n_gb_bin"] + text, count = spy.crossings[0][1], spy.crossings[0][3] + assert json.loads(text)["vecs"]["offs"] == [0] + assert json.loads(text)["rows"][0]["short"] == {"\u0001": ["B", 60 * DIM, 2]} + assert count == 60 * DIM + sum(2 + i % 2 for i in range(60)) + _assert_bit_equal(_stored(db, "Free"), x) + short = _stored(db, "Free", "short") + for i in range(60): + assert ( + np.array(short[i], dtype=np.float32).tobytes() + == x[i][: 2 + i % 2].tobytes() + ) + + +def test_a_float64_or_list_column_keeps_the_text_entry(db, spy): + x = _x(30) + rows = [{"id": i, "embedding": x[i].astype(np.float64)} for i in range(30)] + with db.graph_batch(commit_every=CHUNK) as gb: + gb.create_vertices("Free", rows) + gb.create_vertices( + "Declared", [{"id": i, "embedding": x[i].tolist()} for i in range(30)] + ) + assert spy.entries() == ["arcadedb_n_gb"] * 2 + assert all('"vecs"' not in c[1] for c in spy.crossings) + _assert_bit_equal(_stored(db, "Declared"), x) + + +def test_a_library_without_the_entry_keeps_the_text_path(db, spy, monkeypatch): + monkeypatch.setattr(lib, "GB_BIN", False) + x = _x(30) + with db.graph_batch(commit_every=CHUNK) as gb: + gb.create_vertices("Free", [{"id": i, "embedding": x[i]} for i in range(30)]) + assert spy.entries() == ["arcadedb_n_gb"] + assert '"F"' in spy.crossings[0][1] and '"vecs"' not in spy.crossings[0][1] + _assert_bit_equal(_stored(db, "Free"), x) + + +def test_a_vector_column_outside_the_side_buffer_is_refused(db): + x = _x(4) + with db.graph_batch() as gb: + text = json.dumps( + { + "type": "Free", + "rows": [{"id": i} for i in range(4)], + "vecs": {"names": ["embedding"], "dims": [DIM], "offs": [8]}, + } + ) + buf = x.tobytes() # 4 vectors of 16 floats: offset 8 leaves room for 56, not 64 + reply = json.loads( + lib.call( + "arcadedb_n_gb_bin", + gb._gb, + b"create_vertices", + text.encode(), + buf, + x.size, + ) + ) + assert not reply.get("ok") + assert "outside the side buffer" in json.dumps(reply) + assert db.query("sql", "SELECT count(*) AS c FROM Free").first().get("c") == 0 + + +def test_the_encoder_in_isolation(): + x = _x(5, 4) + rows = [{"id": i} for i in range(5)] + text, buf, count = codec.encode_vertex_args("T", rows, ["v"], [x]) + assert json.loads(text) == { + "type": "T", + "rows": rows, + "vecs": {"names": ["v"], "dims": [4], "offs": [0]}, + } + assert count == 20 and bytes(buf) == x.tobytes() + frozen = x.copy() + frozen.setflags( + write=False + ) # a read-only matrix cannot be viewed writable: its bytes cross instead + assert codec.encode_vertex_args("T", rows, ["v"], [frozen])[1] == x.tobytes() + assert ( + codec._plain_rows(rows) + and codec._plain_rows([{}, {}]) + and codec._plain_rows([]) + ) + assert not codec._plain_rows([{"a": np.float64(1)}]) + assert not codec._plain_rows([{"a": [1]}]) + assert not codec._plain_rows([{1: 2}]) diff --git a/tests/test_vertex_vector_loads.py b/tests/test_vertex_vector_loads.py new file mode 100644 index 00000000..60718b1e --- /dev/null +++ b/tests/test_vertex_vector_loads.py @@ -0,0 +1,344 @@ +"""GraphBatch.create_vertices: float32 vectors of any property cross as one matrix, on both wheels (#337). + +A row value that is a Java ``float[]`` (what ``to_java_float_array`` returns on the JVM wheel) or a numpy float32 array (what it returns on +the native wheel) is a vector candidate for every property, declared ``ARRAY_OF_FLOATS`` or not. Before, the JVM wheel refused a Java +``float[]`` as a candidate and sent the whole load through the property-matrix path (a JPype conversion of every value of every row), and +both wheels wrote one tag per row for a float32 array. Now one float32 matrix per property and chunk crosses (``take_vector_columns``). + +This file pins what is stored (bit for bit, for both wheels, with the one call both wheels share) and the shared split itself. The crossing +counts are in test_graph_batch_vector_crossings_jvm.py (JPype) and test_native_vertex_vector_crossings.py (ctypes). +""" + +import numpy as np +import pytest + +import arcadedb +from arcadedb import type_conversion as tc + +N = 40 +TYPES = ["Declared", "Free"] # embedding declared ARRAY_OF_FLOATS, and not declared +SHAPES = ["to_java_float_array", "float32_ndarray"] + + +def _vectors(dim, n=N, seed=5): + """float32 values a decimal text of 6 digits would not hold: -0.0, the largest, the smallest normal and the smallest subnormal.""" + x = np.random.default_rng(seed).standard_normal((n, dim)).astype(np.float32) + x[0, 0] = np.float32(-0.0) + x[1, 0] = np.finfo(np.float32).max + x[1, 1] = np.finfo(np.float32).tiny + x[1, 2] = np.finfo(np.float32).smallest_subnormal + return x + + +def _make(shape): + if shape == "to_java_float_array": + return ( + arcadedb.to_java_float_array + ) # a float[] on the JVM wheel, a float32 ndarray on the native wheel + return lambda vector: vector + + +def _bits(stored): + return ( + np.array(stored, dtype=np.float64).astype(np.float32).view(np.uint32).tolist() + ) + + +@pytest.fixture +def db(temp_db): + for name in ("Declared", "Bulk", "One"): + temp_db.command("sql", f"CREATE VERTEX TYPE {name}") + temp_db.command("sql", f"CREATE PROPERTY {name}.id LONG") + temp_db.command("sql", "CREATE PROPERTY Declared.embedding ARRAY_OF_FLOATS") + temp_db.command("sql", "CREATE PROPERTY Declared.born DATETIME") + temp_db.command("sql", "CREATE VERTEX TYPE Free") + temp_db.command("sql", "CREATE PROPERTY Free.id LONG") + temp_db.command("sql", "CREATE PROPERTY Free.born DATETIME") + temp_db.command("sql", "CREATE PROPERTY Bulk.embedding ARRAY_OF_FLOATS") + temp_db.command("sql", "CREATE PROPERTY One.embedding ARRAY_OF_FLOATS") + return temp_db + + +def _stored(db, type_name, prop="embedding"): + sql = f"SELECT id, {prop} AS v FROM {type_name}" # nosec B608 - test-owned names + return {r["id"]: r["v"] for r in db.query("sql", sql).to_list()} + + +def _count(db, type_name): + sql = f"SELECT count(*) AS c FROM {type_name}" # nosec B608 - test-owned name + return db.query("sql", sql).first().get("c") + + +@pytest.mark.parametrize("type_name", TYPES) +@pytest.mark.parametrize("dim", [64, 96]) +@pytest.mark.parametrize("shape", SHAPES) +def test_float32_vectors_are_stored_bit_equal(db, type_name, dim, shape): + x = _vectors(dim) + make = _make(shape) + with db.graph_batch() as gb: + rids = gb.create_vertices( + type_name, [{"id": i, "embedding": make(x[i])} for i in range(N)] + ) + assert len(rids) == N and len(set(rids)) == N + stored = _stored(db, type_name) + assert sorted(stored) == list(range(N)) + for i in range(N): + assert _bits(stored[i]) == x[i].view(np.uint32).tolist() + for i in (0, 17, N - 1): # the RIDs come back in row order + sql = f"SELECT id FROM {rids[i]}" # nosec B608 - a RID from this test + assert db.query("sql", sql).first().get("id") == i + + +@pytest.mark.parametrize("shape", SHAPES) +def test_the_matrix_stores_what_the_per_vertex_path_stores(db, shape): + x = _vectors(32, 12) + rows = [{"id": i, "embedding": _make(shape)(x[i])} for i in range(12)] + with db.graph_batch() as gb: + gb.create_vertices("Bulk", rows) + with db.graph_batch() as gb: + for row in rows: + gb.create_vertex("One", **row) + assert _stored(db, "Bulk") == _stored(db, "One") + assert _count(db, "Bulk") == 12 + + +def test_float64_numpy_rows_keep_their_documented_values(db): + x64 = np.random.default_rng(3).random((N, 64)) * 1000.0 + 1e-3 + rows = [{"id": i, "embedding": x64[i]} for i in range(N)] + with db.graph_batch() as gb: + gb.create_vertices("Declared", rows) + gb.create_vertices("Free", rows) + declared, free = _stored(db, "Declared"), _stored(db, "Free") + for i in range(N): + # declared ARRAY_OF_FLOATS: the engine stores floats, so the doubles round to float32 + assert _bits(declared[i]) == x64[i].astype(np.float32).view(np.uint32).tolist() + # undeclared: the list keeps its doubles, exactly + assert np.array(free[i], dtype=np.float64).tolist() == x64[i].tolist() + + +def test_lists_of_floats_keep_their_documented_values(db): + x64 = np.random.default_rng(4).random((N, 40)) * 1000.0 + 1e-3 + rows = [{"id": i, "embedding": x64[i].tolist()} for i in range(N)] + with db.graph_batch() as gb: + gb.create_vertices("Declared", rows) + gb.create_vertices("Free", rows) + declared, free = _stored(db, "Declared"), _stored(db, "Free") + for i in range(N): + assert _bits(declared[i]) == x64[i].astype(np.float32).view(np.uint32).tolist() + assert np.array(free[i], dtype=np.float64).tolist() == x64[i].tolist() + + +@pytest.mark.parametrize("type_name", TYPES) +@pytest.mark.parametrize("shape", SHAPES) +def test_ragged_null_and_missing_vectors_take_the_fallback(db, type_name, shape): + x = _vectors(16, 6) + make = _make(shape) + rows = [ + {"id": 0, "embedding": make(x[0])}, + {"id": 1, "embedding": make(x[1][:9])}, # another length + {"id": 2, "embedding": None}, + {"id": 3}, # no vector + {"id": 4, "embedding": make(x[4])}, + None, # a vertex without properties + ] + with db.graph_batch() as gb: + assert len(gb.create_vertices(type_name, rows)) == 6 + stored = _stored(db, type_name) + assert ( + _count(db, type_name) == 6 and None in stored + ) # the vertex of the None row has no id + assert _bits(stored[0]) == x[0].view(np.uint32).tolist() + assert _bits(stored[1]) == x[1][:9].view(np.uint32).tolist() + assert stored[2] is None and stored[3] is None + assert _bits(stored[4]) == x[4].view(np.uint32).tolist() + + +@pytest.mark.parametrize("type_name", TYPES) +@pytest.mark.parametrize("shape", SHAPES) +def test_a_float64_row_among_float32_rows_loses_nothing(db, type_name, shape): + x = _vectors(16, 8) + make = _make(shape) + rows = [{"id": i, "embedding": make(x[i]), "tag": f"t{i}"} for i in range(8)] + rows[5] = {"id": 5, "tag": "t5", "embedding": make(x[5].astype(np.float64))} + with db.graph_batch() as gb: + gb.create_vertices(type_name, rows) + stored, tags = _stored(db, type_name), _stored(db, type_name, "tag") + assert [tags[i] for i in range(8)] == [f"t{i}" for i in range(8)] + for i in range(8): + if type_name == "Free" and i == 5 and shape == "float32_ndarray": + # a float64 array on an undeclared property keeps its doubles (it is no float32 vector) + assert ( + np.array(stored[i], dtype=np.float64).tolist() + == x[5].astype(np.float64).tolist() + ) + else: + assert _bits(stored[i]) == x[i].view(np.uint32).tolist() + + +@pytest.mark.parametrize("type_name", TYPES) +@pytest.mark.parametrize("shape", SHAPES) +def test_the_other_properties_stay_with_their_rows(db, type_name, shape): + x = _vectors(16, 12) + make = _make(shape) + rows = [] + for i in range(12): + row = { + "id": i, + "name": f"n{i}", + "score": i * 0.5, + "flag": i % 2 == 0, + "note": None, + "embedding": make(x[i]), + } + if i % 4 == 3: # another key order: the rows no longer share one + row = dict(reversed(list(row.items()))) + rows.append(row) + with db.graph_batch() as gb: + gb.create_vertices(type_name, rows) + sql = f"SELECT id, name, score, flag, embedding FROM {type_name} ORDER BY id" # nosec B608 - test-owned name + got = db.query("sql", sql).to_list() + assert [r["id"] for r in got] == list(range(12)) + for i, r in enumerate(got): + assert ( + r["name"] == f"n{i}" and r["score"] == i * 0.5 and r["flag"] is (i % 2 == 0) + ) + assert _bits(r["embedding"]) == x[i].view(np.uint32).tolist() + + +@pytest.mark.parametrize("type_name", TYPES) +@pytest.mark.parametrize("shape", SHAPES) +def test_two_vector_properties_and_a_date_time_in_one_load(db, type_name, shape): + from datetime import datetime + + x, y = _vectors(8, 20), _vectors(8, 20, seed=9) + make = _make(shape) + rows = [ + { + "id": i, + "embedding": make(x[i]), + "second": make(y[i]), + "born": datetime(2020, 1, 2, 3, 4, 5), + } + for i in range(20) + ] + with db.graph_batch() as gb: + gb.create_vertices(type_name, rows) + first, second = _stored(db, type_name), _stored(db, type_name, "second") + for i in range(20): + assert _bits(first[i]) == x[i].view(np.uint32).tolist() + assert _bits(second[i]) == y[i].view(np.uint32).tolist() + sql = f"SELECT born FROM {type_name} WHERE id = 7" # nosec B608 - test-owned name + assert db.query("sql", sql).first().get("born") == datetime(2020, 1, 2, 3, 4, 5) + + +@pytest.mark.parametrize("type_name", TYPES) +@pytest.mark.parametrize("shape", SHAPES) +def test_a_load_over_the_chunk_boundary(db, type_name, shape): + make = _make(shape) + with db.graph_batch() as gb: + # the rows per crossing: the JVM wheel's 100,000, the native wheel's commit_every (1000 unless graph_batch() got another value) + chunk = getattr(gb, "_vertex_chunk", None) or gb._BULK_CHUNK + n = chunk + 3 + x = _vectors(4, n) + rids = gb.create_vertices( + type_name, [{"id": i, "embedding": make(x[i])} for i in range(n)] + ) + assert len(rids) == n and len(set(rids)) == n + assert _count(db, type_name) == n + stored = _stored(db, type_name) + got = np.array([stored[i] for i in range(n)], dtype=np.float64).astype(np.float32) + assert (got.view(np.uint32) == x.view(np.uint32)).all() + for i in ( + 0, + chunk - 1, + chunk, + n - 1, + ): # the RIDs come back in row order across the boundary + sql = f"SELECT id FROM {rids[i]}" # nosec B608 - a RID from this test + assert db.query("sql", sql).first().get("id") == i + + +@pytest.mark.parametrize("shape", SHAPES) +def test_insert_many_takes_the_same_vectors(db, shape): + db.command("sql", "CREATE DOCUMENT TYPE Doc") + db.command("sql", "CREATE PROPERTY Doc.id LONG") + db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + x = _vectors(32, N) + make = _make(shape) + assert ( + db.insert_many("Doc", [{"id": i, "embedding": make(x[i])} for i in range(N)]) + == N + ) + stored = _stored(db, "Doc") + for i in range(N): + assert _bits(stored[i]) == x[i].view(np.uint32).tolist() + + +# ---- the split itself (no database): the one function both wheels call + + +def _rows(x, **extra): + return [{"a": i, "b": f"s{i}", "v": x[i], **extra} for i in range(len(x))] + + +def test_take_vector_columns_takes_a_float32_column_of_any_property(): + x = _vectors(5, 7) + names, matrices, rest = tc.take_vector_columns(_rows(x), (), undeclared=True) + assert names == ["v"] and matrices[0].dtype == np.float32 + assert matrices[0].shape == (7, 5) and matrices[0].flags.c_contiguous + assert matrices[0].tobytes() == x.tobytes() # bit for bit + assert rest == [{"a": i, "b": f"s{i}"} for i in range(7)] + assert all(list(row) == ["a", "b"] for row in rest) + + +def test_take_vector_columns_leaves_an_undeclared_column_that_is_no_float32_vector(): + x = _vectors(5, 7) + as_lists = [{"v": row.tolist()} for row in x] + as_doubles = [{"v": row.astype(np.float64)} for row in x] + assert tc.take_vector_columns(as_lists, (), undeclared=True) is None + assert tc.take_vector_columns(as_doubles, (), undeclared=True) is None + assert tc.take_vector_columns(as_lists, ["v"])[1][0].shape == (7, 5) + assert tc.take_vector_columns(as_doubles, ["v"])[1][0].shape == (7, 5) + assert ( + tc.take_vector_columns(_rows(x), ()) is None + ) # undeclared names are tried only when asked + + +def test_take_vector_columns_declines_a_ragged_null_or_missing_column(): + x = _vectors(5, 4) + for rows in ( + [{"v": x[0]}, {"v": x[1][:3]}], + [{"v": x[0]}, {"v": None}], + [{"v": x[0]}, {"w": 1}], + [{"v": x[0]}, None], + [{"v": np.zeros(0, dtype=np.float32)}], + [{"v": x}], # a 2-D array per row + ): + assert tc.take_vector_columns(rows, ["v"], undeclared=True) is None + + +def test_take_vector_columns_rebuilds_the_rest_rows_in_each_rows_own_key_order(): + x = _vectors(3, 2) + rows = [{"a": 1, "v": x[0], "b": 2}, {"b": 3, "v": x[1], "a": 4}] + rest = tc.take_vector_columns(rows, (), undeclared=True)[2] + assert rest == [{"a": 1, "b": 2}, {"b": 3, "a": 4}] + assert [list(row) for row in rest] == [["a", "b"], ["b", "a"]] + only = tc.take_vector_columns([{"v": x[0]}, {"v": x[1]}], (), undeclared=True)[2] + assert only == [{}, {}] and only[0] is not only[1] + two = [{"x": r, "y": r * 2, "k": i} for i, r in enumerate(x)] + names, _, rest = tc.take_vector_columns(two, (), undeclared=True) + assert sorted(names) == ["x", "y"] and rest == [{"k": 0}, {"k": 1}] + + +def test_take_vector_columns_takes_a_strided_or_read_only_row(): + wide = _vectors(10, 6) + strided = tc.take_vector_columns( + [{"v": wide[i, ::2]} for i in range(6)], (), undeclared=True + )[1][0] + assert strided.tobytes() == np.ascontiguousarray(wide[:, ::2]).tobytes() + frozen = wide.copy() + frozen.setflags(write=False) + matrix = tc.take_vector_columns( + [{"v": frozen[i]} for i in range(6)], (), undeclared=True + )[1][0] + assert matrix.tobytes() == wide.tobytes() and matrix.flags.writeable