From 0a2e681080086b83603445d9905419148736061e Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Sat, 24 Oct 2026 21:22:09 +0900 Subject: [PATCH] Bridge: float vectors cross as one binary buffer, not as JSON text or one JPype call per element A vector load spent most of its time on text: Python wrote every float as a decimal number and the engine parsed it back (insert_many), or JPype converted every list element on its own (create_vertices, Document.set). For a property declared ARRAY_OF_FLOATS the engine converts each number of a list to float anyway, so the vectors of a chunk now cross as one float32 matrix (FloatVectors.Matrix, one buffer copy) beside the JSON, and each document gets its own float[] row Java-side. Properties that are not declared ARRAY_OF_FLOATS, and vectors that are ragged, missing or null, keep the exact JSON path. Also: insert_columns accepts a 2-D float32 array as a column of vectors; a numpy array is accepted as a property value everywhere (Document.set, insert_many, create_vertices; it raised TypeError before); the allowed_rids filter of a vector search builds its RID set in one crossing (GraphCalls.ridSet) instead of two per RID. Tests: tests/test_float_vector_bulk.py (39 tests: exact float32 equality with the engine's own list conversion on doubles a float32 cannot hold, which values keep the JSON path, every entry point, parallel mode); tests/test_vector.py rid-conversion test now passes a Java RID object so the list takes the per-RID path; tests/test_python_callbacks.py numpy return value test follows the new conversion. Ported from humemai/arcadedb-embedded-python#496. The user docs and the speed table generated there live in that repository's docs; this repository has no user docs yet. Co-Authored-By: Claude Sonnet 5.5 --- src/arcadedb_embedded/core.py | 71 +++- src/arcadedb_embedded/graph_batch.py | 69 +++- src/arcadedb_embedded/type_conversion.py | 104 +++++ src/arcadedb_embedded/vector.py | 11 +- .../com/arcadedb/python/DocumentBatcher.java | 45 ++- .../com/arcadedb/python/FloatVectors.java | 83 ++++ src/java/com/arcadedb/python/GraphCalls.java | 13 + .../com/arcadedb/python/VertexBatcher.java | 18 +- tests/test_float_vector_bulk.py | 370 ++++++++++++++++++ tests/test_python_callbacks.py | 15 +- tests/test_vector.py | 10 +- 11 files changed, 766 insertions(+), 43 deletions(-) create mode 100644 src/java/com/arcadedb/python/FloatVectors.java create mode 100644 tests/test_float_vector_bulk.py diff --git a/src/arcadedb_embedded/core.py b/src/arcadedb_embedded/core.py index 63f8a43..4e781af 100644 --- a/src/arcadedb_embedded/core.py +++ b/src/arcadedb_embedded/core.py @@ -29,7 +29,12 @@ from .jvm import start_jvm from .results import ResultSet from .transactions import TransactionContext -from .type_conversion import _is_numpy_bool, convert_python_to_java, json_bulk_dumps +from .type_conversion import ( + _is_numpy_bool, + convert_python_to_java, + json_bulk_dumps, + split_float_vectors, +) from .vector import to_java_float_array _LOGGER = get_logger(__name__) @@ -184,7 +189,8 @@ def _column_to_java(name, values): and any other sequence convert per element into an ``Object[]`` (a ``str`` reuses one Java String per distinct value, ``None`` is a null). A pandas ``Series`` is read through ``to_numpy``, with a nullable dtype's ```` - as ``None``. Kinds that do not cross natively raise ``TypeError`` rather + as ``None``. A 2-D ``float32`` array is a column of vectors: it crosses as + one ``float[]`` and each document gets one row of it. Kinds that do not cross natively raise ``TypeError`` rather than being stored as something else. """ if hasattr(values, "to_numpy") and not isinstance(values, (list, tuple)): @@ -194,9 +200,26 @@ def _column_to_java(name, values): else: # a pandas extension dtype (Int64, string, boolean, category): its is a null values = values.to_numpy(dtype=object, na_value=None) if _np is not None and isinstance(values, _np.ndarray): + if values.ndim == 2 and values.dtype == _np.float32 and values.shape[1] > 0: + # a column of vectors: ONE float[] holds the whole matrix and the + # documents get a row each, Java-side (FloatVectors.Matrix) + return ( + _java_class("com.arcadedb.python.FloatVectors$Matrix")( + jpype.JArray(jpype.JFloat)( + _np.ascontiguousarray(values).reshape(-1) + ), + int(values.shape[1]), + ), + int(values.shape[0]), + ) if values.ndim != 1: + hint = ( + "; a 2-D float32 matrix is accepted as a column of vectors" + if values.ndim == 2 + else "" + ) raise ValueError( - f"column {name!r} is {values.ndim}-dimensional; a column is one-dimensional" + f"column {name!r} is {values.ndim}-dimensional; a column is one-dimensional{hint}" ) kind = values.dtype.kind if kind in "iu": @@ -664,7 +687,13 @@ def insert_many( generator of any length can be loaded (#294). Values must be JSON-representable (str/int/float/bool/None and nested lists/dicts); a chunk containing other types (e.g. datetime, bytes) falls back - transparently to the per-row path. + transparently to the per-row path. A property declared + ``ARRAY_OF_FLOATS`` whose value in every row of a chunk is a list or a + numpy array of numbers of one length crosses as one binary float32 + matrix beside the JSON instead of as text (``FloatVectors``); the stored + values are the list's, since the engine converts a list for such a + property to float. A vector that is ragged, missing or null in a row + stays in the JSON. Args: type_name: Target document type (must exist). @@ -770,8 +799,11 @@ def insert_many( def _insert_many_chunk(self, type_name, chunk): """Insert one chunk of insert_many rows inside the open transaction.""" + # Float vectors for ARRAY_OF_FLOATS properties cross as one binary matrix + # beside the JSON, not as text (see split_float_vectors). + split = split_float_vectors(self._java_db, type_name, chunk) try: - payload = json_bulk_dumps(chunk) + payload = json_bulk_dumps(split[0] if split else chunk) except (TypeError, ValueError): # Values the JSON text cannot carry unchanged (numpy integer # scalars, which json.dumps rejects; an integer beyond 64 bits, @@ -787,9 +819,13 @@ def _insert_many_chunk(self, type_name, chunk): try: # commitEvery 0 inside a transaction that is already open: the # helper only inserts, and the transaction stays with insert_many. - _java_class("com.arcadedb.python.DocumentBatcher").insertManyJson( - self._java_db, type_name, payload, 0, False - ) + batcher = _java_class("com.arcadedb.python.DocumentBatcher") + if split: + batcher.insertManyJson( + self._java_db, type_name, payload, 0, False, split[1], split[2] + ) + else: + batcher.insertManyJson(self._java_db, type_name, payload, 0, False) except Exception as e: raise ArcadeDBError(f"Failed to bulk-insert into '{type_name}': {e}") from e @@ -801,8 +837,9 @@ def _insert_many_parallel(self, type_name, first, rows): chunk = first try: while chunk: + split = split_float_vectors(self._java_db, type_name, chunk) try: - payload = json_bulk_dumps(chunk) + payload = json_bulk_dumps(split[0] if split else chunk) except (TypeError, ValueError): payload = None if payload is None: @@ -827,9 +864,14 @@ def _insert_many_parallel(self, type_name, first, rows): raise else: try: - failures = batcher.insertManyJsonParallel( - self._java_db, type_name, payload - ) + if split: + failures = batcher.insertManyJsonParallel( + self._java_db, type_name, payload, split[1], split[2] + ) + else: + failures = batcher.insertManyJsonParallel( + self._java_db, type_name, payload + ) except Exception as e: raise self._parallel_load_failed(type_name, e) from e # The writers report a rejected record only through its @@ -887,7 +929,10 @@ def insert_columns( engine): 8.2 s against 18.4 s for ``insert_many`` (2.24x); every arm stored the same sums and count. Use it whenever the data already lives in columns (a pandas ``DataFrame``, a parquet batch, numpy arrays); - ``insert_many`` stays the path for a list of dicts. + ``insert_many`` stays the path for a list of dicts. A 2-D ``float32`` + array is a column of vectors: the whole matrix crosses as one + ``float[]`` and each document gets one row of it (an ``ARRAY_OF_FLOATS`` + property). Args: type_name: Target document type (must exist). diff --git a/src/arcadedb_embedded/graph_batch.py b/src/arcadedb_embedded/graph_batch.py index 9d8b3e3..1e1f45f 100644 --- a/src/arcadedb_embedded/graph_batch.py +++ b/src/arcadedb_embedded/graph_batch.py @@ -11,7 +11,13 @@ from ._logging import get_logger, log_swallowed_exception from .exceptions import ArcadeDBError from .graph import Document, Vertex -from .type_conversion import convert_python_to_java, json_bulk_scalar_ok +from .type_conversion import ( + convert_python_to_java, + declared_float_vectors, + is_vector_candidate, + json_bulk_scalar_ok, + split_float_vectors, +) _LOGGER = get_logger(__name__) @@ -180,7 +186,10 @@ def create_vertices( type_name: Vertex type name. 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. + 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. """ self._check_not_closed() # From engine 26.10.1 the engine refuses to run inside a transaction the @@ -228,23 +237,31 @@ 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). Returns None when unavailable/unsuitable so the caller falls - back to the property-matrix path.""" + 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). + Returns None when unavailable/unsuitable so the caller falls back to + the property-matrix path.""" try: vertex_batcher = jpype.JClass("com.arcadedb.python.VertexBatcher") + declared = set(declared_float_vectors(self._java_db, type_name)) except Exception: return None for row in rows: if row: for key, value in row.items(): - if type(key) is not str or not json_bulk_scalar_ok(value): + 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): # e.g. datetime/bytes (the matrix path preserves types), or a # value JSON would change: an integer beyond 64 bits, NaN or - # Infinity, a lone surrogate; or 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 + # Infinity, a lone surrogate return None import json @@ -253,17 +270,39 @@ def _create_vertices_json_bulk(self, type_name, rows): rids = [] for start in range(0, len(rows), self._BULK_CHUNK): chunk = rows[start : start + self._BULK_CHUNK] - joined = str( - vertex_batcher.createVerticesJson( - self._java_graph_batch, - type_name, - json.dumps([row or {} for row in chunk]), + split = split_float_vectors(self._java_db, type_name, chunk, declared) + rest = split[0] if split else chunk + if declared and not self._json_bulk_ok(rest): + # a list the matrix cannot carry (ragged, a null, not numbers) + # is left in this chunk: the property-matrix path takes the chunk + java_rids = self._java_graph_batch.createVertices( + type_name, self._to_java_property_matrix(chunk) ) - ) + rids.extend(str(rid) for rid in java_rids) + continue + text = json.dumps([row or {} for row in rest]) + if split: + joined = vertex_batcher.createVerticesJson( + self._java_graph_batch, type_name, text, split[1], split[2] + ) + else: + joined = vertex_batcher.createVerticesJson( + self._java_graph_batch, type_name, text + ) + joined = str(joined) if joined: 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 + def new_edge( self, source_vertex_or_rid, diff --git a/src/arcadedb_embedded/type_conversion.py b/src/arcadedb_embedded/type_conversion.py index 1e4fce0..b552319 100644 --- a/src/arcadedb_embedded/type_conversion.py +++ b/src/arcadedb_embedded/type_conversion.py @@ -13,6 +13,11 @@ import jpype +try: + import numpy as _np +except ImportError: # numpy is optional + _np = None + class _JavaTimeTypes(NamedTuple): instant: Any @@ -506,6 +511,16 @@ def convert_python_to_java(value: Any) -> Any: # Not a bool subclass: JPype would store it as the Double 1.0 or 0.0. return bool(value) + if _np is not None and isinstance(value, _np.ndarray): + # A numpy array reached the JVM as a TypeError ("only integer scalar + # arrays can be converted to a scalar index") from JPype. A float32 vector + # crosses as ONE float[] (one buffer copy), what an ARRAY_OF_FLOATS + # property stores; any other array is the list it would have been + # (exact values: a float64 array keeps its doubles). + if value.ndim == 1 and value.dtype == _np.float32: + return jpype.JArray(jpype.JFloat)(value) + return convert_python_to_java(value.tolist()) + java_python_types = _get_java_python_types() if isinstance(value, Decimal): @@ -707,3 +722,92 @@ def json_bulk_dumps(rows: Any) -> str: value ) # nested lists and dicts, and subclasses of the types above return json.dumps(rows, allow_nan=False) + + +_FLOAT_VECTORS = None # the bridge's FloatVectors class once resolved; False when the bridge jar lacks it + + +def _float_vectors_class(): + global _FLOAT_VECTORS + if _FLOAT_VECTORS is None: + try: + _FLOAT_VECTORS = jpype.JClass("com.arcadedb.python.FloatVectors") + except Exception: # an older bridge jar: the JSON text carries the vectors + _FLOAT_VECTORS = False + return _FLOAT_VECTORS or None + + +def declared_float_vectors(java_db, type_name): + """The names of the ARRAY_OF_FLOATS properties of `type_name` (inherited ones too), [] when there are none or no numpy.""" + vectors = _float_vectors_class() if _np is not None else None + if vectors is None: + return [] + return [str(name) for name in vectors.declared(java_db, type_name)] + + +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)) + + +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. + + 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). + """ + 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 = [] + matrices = [] + for name in declared: + try: + column = [row[name] for row in rows] + except (KeyError, TypeError): + continue + if not is_vector_candidate(column[0]): + continue + try: + array = _np.array(column) + except (ValueError, TypeError): + continue + if array.ndim != 2 or array.shape[1] == 0 or array.dtype.kind not in "fiu": + 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) + taken.append(name) + matrices.append( + vectors.Matrix(jpype.JArray(jpype.JFloat)(flat), int(array.shape[1])) + ) + 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 ( + rest, + jpype.JArray(jpype.JString)(taken), + jpype.JArray(jpype.JObject)(matrices), + ) diff --git a/src/arcadedb_embedded/vector.py b/src/arcadedb_embedded/vector.py index 717553f..0590ee7 100644 --- a/src/arcadedb_embedded/vector.py +++ b/src/arcadedb_embedded/vector.py @@ -386,10 +386,19 @@ def _build_allowed_rids_set(self, allowed_rids): if not allowed_rids: return None + from .graph import _graph_calls + + values = list(allowed_rids) + calls = _graph_calls() + if calls is not None and all(type(v) is str for v in values): + # RIDs given as text: one crossing builds the whole set (a RID + # constructor and an add() call per RID took about 5 us each) + return calls.ridSet(jpype.JArray(jpype.JString)(values)) + HashSet = jpype.JClass("java.util.HashSet") allowed_rids_set = HashSet() - for rid_value in allowed_rids: + for rid_value in values: allowed_rids_set.add(self._database.to_java_rid(rid_value)) return allowed_rids_set diff --git a/src/java/com/arcadedb/python/DocumentBatcher.java b/src/java/com/arcadedb/python/DocumentBatcher.java index b77d62e..848f18f 100644 --- a/src/java/com/arcadedb/python/DocumentBatcher.java +++ b/src/java/com/arcadedb/python/DocumentBatcher.java @@ -11,7 +11,8 @@ * the async executor's parallel bucket writers (insertManyJsonParallel; the * Python insert_many wrapper waits for completion itself, then reads the * failures the writers reported). insertColumns takes whole columns instead of - * JSON (Database.insert_columns). The boxDoubles/boxLongs + * JSON (Database.insert_columns). Float vectors for ARRAY_OF_FLOATS properties cross as binary matrices (FloatVectors), beside + * the JSON or as a column. The boxDoubles/boxLongs * helpers below serve AsyncExecutor.append_samples' numpy fast path. * * JSON-representable property values only (str/int/float/bool/null and @@ -37,10 +38,21 @@ private DocumentBatcher() { public static long insertManyJson(final Database db, final String typeName, final String jsonRows, final int commitEvery, final boolean parallel) { + return insertManyJson(db, typeName, jsonRows, commitEvery, parallel, null, null); + } + + /** + * As {@link #insertManyJson(Database, String, String, int, boolean)}, with float vectors that did not go through the JSON + * text: {@code vectorNames[c]} is set on row i to row i of {@code vectors[c]}, a {@link FloatVectors.Matrix} of exactly as + * many rows as the JSON has. Both are null when there are none. + */ + public static long insertManyJson(final Database db, final String typeName, final String jsonRows, + final int commitEvery, final boolean parallel, final String[] vectorNames, final Object[] vectors) { final JSONArray rows = new JSONArray(jsonRows); final int n = rows.length(); + FloatVectors.check(vectorNames, vectors, n); if (parallel) { - insertManyJsonParallel(db, typeName, rows); + insertManyJsonParallel(db, typeName, rows, vectorNames, vectors); return n; } // A caller's own transaction is the caller's to commit or roll back: batch @@ -52,7 +64,7 @@ public static long insertManyJson(final Database db, final String typeName, fina try { for (int i = 0; i < n; i++) { final MutableDocument doc = db.newDocument(typeName); - fill(doc, rows.getJSONObject(i)); + fill(doc, rows.getJSONObject(i), i, vectorNames, vectors); doc.save(); if (!wasActive && commitEvery > 0 && (i + 1) % commitEvery == 0) { db.commit(); @@ -150,9 +162,11 @@ else if (col instanceof boolean[] a) len = a.length; else if (col instanceof Object[] a) len = a.length; + else if (col instanceof FloatVectors.Matrix m) + len = m.rows(); else throw new IllegalArgumentException("column '" + names[c] + "' is a " + (col == null ? "null" : col.getClass().getName()) - + ", not a long[], double[], boolean[] or Object[]"); + + ", not a long[], double[], boolean[], Object[] or FloatVectors.Matrix"); if (len != n) throw new IllegalArgumentException("column '" + names[c] + "' has " + len + " values for " + n + " rows"); } @@ -168,6 +182,8 @@ else if (col instanceof double[] a) v = a[i]; else if (col instanceof boolean[] a) v = a[i]; + else if (col instanceof FloatVectors.Matrix m) + v = m.row(i); else v = ((Object[]) col)[i]; doc.set(names[c], v); @@ -201,16 +217,25 @@ public String getFirstMessage() { } public static AsyncFailures insertManyJsonParallel(final Database db, final String typeName, final String jsonRows) { - return insertManyJsonParallel(db, typeName, new JSONArray(jsonRows)); + return insertManyJsonParallel(db, typeName, jsonRows, null, null); + } + + /** The parallel twin of {@link #insertManyJson(Database, String, String, int, boolean, String[], Object[])}. */ + public static AsyncFailures insertManyJsonParallel(final Database db, final String typeName, final String jsonRows, + final String[] vectorNames, final Object[] vectors) { + final JSONArray rows = new JSONArray(jsonRows); + FloatVectors.check(vectorNames, vectors, rows.length()); + return insertManyJsonParallel(db, typeName, rows, vectorNames, vectors); } - private static AsyncFailures insertManyJsonParallel(final Database db, final String typeName, final JSONArray rows) { + private static AsyncFailures insertManyJsonParallel(final Database db, final String typeName, final JSONArray rows, + final String[] vectorNames, final Object[] vectors) { final AsyncFailures failures = new AsyncFailures(); final ErrorCallback onError = failures::record; final int n = rows.length(); for (int i = 0; i < n; i++) { final MutableDocument doc = db.newDocument(typeName); - fill(doc, rows.getJSONObject(i)); + fill(doc, rows.getJSONObject(i), i, vectorNames, vectors); db.async().createRecord(doc, null, onError); } return failures; @@ -233,10 +258,14 @@ public static Object[] boxLongs(final long[] a) { return out; } - private static void fill(final MutableDocument doc, final JSONObject row) { + private static void fill(final MutableDocument doc, final JSONObject row, final int i, final String[] vectorNames, + final Object[] vectors) { // toMap() turns nested objects and arrays into plain Map and List. Storing the parsed JSONArray itself reads back // as a JSONArray inside the transaction that wrote it, and only becomes a list once the record is serialized. for (final Map.Entry entry : row.toMap().entrySet()) doc.set(entry.getKey(), entry.getValue()); + if (vectorNames != null) + for (int c = 0; c < vectorNames.length; c++) + doc.set(vectorNames[c], ((FloatVectors.Matrix) vectors[c]).row(i)); } } diff --git a/src/java/com/arcadedb/python/FloatVectors.java b/src/java/com/arcadedb/python/FloatVectors.java new file mode 100644 index 0000000..74a93f3 --- /dev/null +++ b/src/java/com/arcadedb/python/FloatVectors.java @@ -0,0 +1,83 @@ +/* + * Python-bindings bridge: float vectors that cross as one binary buffer instead of JSON text or one JPype call per element. + * + * Writing a 128-dimension embedding as JSON costs Python about 450 ns per number (float to text) and the engine's JSON parser + * about the same again (text to Double, boxed in a List, converted back to float[] by the engine), so a vector load spent + * most of its time on the text. Here a whole column of vectors crosses as ONE float[] (a numpy buffer copy) plus the + * dimension, and each document gets its own float[] slice, which is what the engine stores for an ARRAY_OF_FLOATS property. + * + * Python uses this only for properties declared ARRAY_OF_FLOATS (declared() lists them): the engine converts every number of a + * list for such a property to float, so a float32 buffer stores the same values as the list did. Any other property keeps the + * JSON path, which keeps its exact values (Long and Double stay what they were). + * + * Compiled into arcadedb-python-bridge.jar during the wheel build and consumed by Database.insert_many(), Database.insert_columns() + * and GraphBatch.create_vertices(). + */ +package com.arcadedb.python; + +import com.arcadedb.database.Database; +import com.arcadedb.schema.DocumentType; +import com.arcadedb.schema.Property; +import com.arcadedb.schema.Type; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +public final class FloatVectors { + + private FloatVectors() { + } + + /** The names of the properties of {@code typeName}, or of a type it extends, that are declared ARRAY_OF_FLOATS. */ + public static String[] declared(final Database db, final String typeName) { + if (!db.getSchema().existsType(typeName)) + return new String[0]; + final DocumentType type = db.getSchema().getType(typeName); + final List names = new ArrayList<>(); + for (final String name : type.getPolymorphicPropertyNames()) { + final Property property = type.getPolymorphicPropertyIfExists(name); + if (property != null && property.getType() == Type.ARRAY_OF_FLOATS) + names.add(name); + } + return names.toArray(new String[0]); + } + + /** {@code n} vectors of {@code dim} floats each, back to back in one array. */ + public static final class Matrix { + private final float[] data; + private final int dim; + + public Matrix(final float[] data, final int dim) { + if (dim < 1) + throw new IllegalArgumentException("a vector needs at least one dimension, got " + dim); + if (data.length % dim != 0) + throw new IllegalArgumentException(data.length + " floats are not a whole number of " + dim + "-dimension vectors"); + this.data = data; + this.dim = dim; + } + + public int rows() { + return data.length / dim; + } + + /** A copy of row {@code i}: the document owns its array, whatever happens to the buffer afterwards. */ + public float[] row(final int i) { + return Arrays.copyOfRange(data, i * dim, (i + 1) * dim); + } + } + + /** Every matrix holds exactly {@code n} rows, and there is one name for each matrix. */ + static void check(final String[] names, final Object[] vectors, final int n) { + if (names == null && vectors == null) + return; + if (names == null || vectors == null || names.length != vectors.length) + throw new IllegalArgumentException("vector names and matrices differ in number"); + for (int c = 0; c < names.length; c++) { + if (!(vectors[c] instanceof Matrix m)) + throw new IllegalArgumentException("'" + names[c] + "' is not a float matrix"); + if (m.rows() != n) + throw new IllegalArgumentException("'" + names[c] + "' has " + m.rows() + " vectors for " + n + " rows"); + } + } +} diff --git a/src/java/com/arcadedb/python/GraphCalls.java b/src/java/com/arcadedb/python/GraphCalls.java index 0b388c4..6ff6838 100644 --- a/src/java/com/arcadedb/python/GraphCalls.java +++ b/src/java/com/arcadedb/python/GraphCalls.java @@ -19,7 +19,9 @@ import com.arcadedb.utility.Pair; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Set; public final class GraphCalls { @@ -48,6 +50,17 @@ public static Edge[] bothEdges(final Vertex vertex, final String... labels) { return drain(vertex.getEdges(Vertex.DIRECTION.BOTH, labels)); } + /** + * A set of RIDs from their text, for the {@code allowed_rids} filter of a vector search: building it from Python took a RID + * constructor call and an {@code add} call per RID (about 5 microseconds each), here the strings cross as one array. + */ + public static Set ridSet(final String[] rids) { + final Set set = new HashSet<>(Math.max(16, (int) (rids.length / 0.75f) + 1)); + for (final String rid : rids) + set.add(new RID(rid)); + return set; + } + /** * The records and scores of vector search hits: for each {@code Pair}, {@code db.lookupByRID(rid, true)} * and the score as a double. Returns {@code {Object[] records, double[] scores}}; a record is null where the engine diff --git a/src/java/com/arcadedb/python/VertexBatcher.java b/src/java/com/arcadedb/python/VertexBatcher.java index 2c93563..5c52998 100644 --- a/src/java/com/arcadedb/python/VertexBatcher.java +++ b/src/java/com/arcadedb/python/VertexBatcher.java @@ -24,16 +24,32 @@ private VertexBatcher() { } public static String createVerticesJson(final GraphBatch batch, final String typeName, final String jsonRows) { + return createVerticesJson(batch, typeName, jsonRows, null, null); + } + + /** + * As {@link #createVerticesJson(GraphBatch, String, String)}, with float vectors that did not go through the JSON text: + * {@code vectorNames[c]} is set on vertex i to row i of {@code vectors[c]}, a {@link FloatVectors.Matrix} with one row per + * JSON row. Both are null when there are none. + */ + public static String createVerticesJson(final GraphBatch batch, final String typeName, final String jsonRows, + final String[] vectorNames, final Object[] vectors) { final JSONArray rows = new JSONArray(jsonRows); + FloatVectors.check(vectorNames, vectors, rows.length()); + final int extra = vectorNames == null ? 0 : vectorNames.length; final Object[][] matrix = new Object[rows.length()][]; for (int i = 0; i < rows.length(); i++) { final JSONObject row = rows.getJSONObject(i); - final Object[] flat = new Object[row.length() * 2]; + final Object[] flat = new Object[(row.length() + extra) * 2]; int j = 0; for (final String key : row.keySet()) { flat[j++] = key; flat[j++] = row.isNull(key) ? null : row.get(key); } + for (int c = 0; c < extra; c++) { + flat[j++] = vectorNames[c]; + flat[j++] = ((FloatVectors.Matrix) vectors[c]).row(i); + } matrix[i] = flat; } final RID[] rids = batch.createVertices(typeName, matrix); diff --git a/tests/test_float_vector_bulk.py b/tests/test_float_vector_bulk.py new file mode 100644 index 0000000..c5f6a69 --- /dev/null +++ b/tests/test_float_vector_bulk.py @@ -0,0 +1,370 @@ +"""Float vectors cross the bridge as one binary matrix, not as JSON text or one JPype call per element. + +A list of 128 floats written as JSON text cost Python about 450 ns a number and the engine's parser about the same again, so +a vector load spent most of its time on the text (bridge class ``FloatVectors``). For a property declared ARRAY_OF_FLOATS the +engine converts every number of a list to float, so a float32 matrix stores the same values as the list did: these tests pin +that equality with doubles a float32 cannot hold, pin which values keep the JSON path (exact doubles on an undeclared property, +ragged or null vectors), and cover every entry point: ``insert_many`` (also parallel), ``insert_columns`` with a 2-D float32 +column, ``GraphBatch.create_vertices``, ``Document.set`` with a numpy array, and the ``allowed_rids`` filter of a vector search. +""" + +from datetime import datetime + +import numpy as np +import pytest + +import arcadedb_embedded as arcadedb +from arcadedb_embedded import type_conversion + +DIM = 8 +N = 50 + + +@pytest.fixture +def vec_db(temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Doc") + temp_db.command("sql", "CREATE PROPERTY Doc.id LONG") + temp_db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + temp_db.command( + "sql", "CREATE DOCUMENT TYPE Plain" + ) # embedding undeclared: the list keeps its doubles + temp_db.command("sql", "CREATE VERTEX TYPE Vx") + temp_db.command("sql", "CREATE PROPERTY Vx.id LONG") + temp_db.command("sql", "CREATE PROPERTY Vx.embedding ARRAY_OF_FLOATS") + temp_db.command("sql", "CREATE DOCUMENT TYPE Parent") + temp_db.command("sql", "CREATE PROPERTY Parent.embedding ARRAY_OF_FLOATS") + temp_db.command("sql", "CREATE DOCUMENT TYPE Child EXTENDS Parent") + return temp_db + + +def _doubles(n=N, dim=DIM, seed=7): + """float64 values a float32 cannot hold, so a path that kept doubles or lost bits shows.""" + return np.random.default_rng(seed).random((n, dim)) * 1000.0 + 1e-3 + + +def _rows(x): + return [{"id": i, "embedding": x[i].tolist()} for i in range(len(x))] + + +def _stored(db, type_name): + """{id: stored vector as float64} for every record of the type.""" + out = {} + query = ( + f"SELECT id, embedding FROM {type_name}" # nosec B608 - test-owned type name + ) + for r in db.query("sql", query).to_list(): + out[r["id"]] = np.array(r["embedding"], dtype=np.float64) + return out + + +def _assert_float32_of(stored, x): + assert sorted(stored) == list(range(len(x))) + for i, vec in stored.items(): + assert np.array_equal(vec, x[i].astype(np.float32).astype(np.float64)), i + + +class TestDecision: + """split_float_vectors takes a property out of the JSON only when the float32 matrix stores the same values.""" + + def _split(self, db, type_name, rows): + return type_conversion.split_float_vectors( + db.get_java_database(), type_name, rows + ) + + def test_a_declared_vector_column_is_taken_out(self, vec_db): + rows = _rows(_doubles()) + rest, names, matrices = self._split(vec_db, "Doc", rows) + assert [str(n) for n in names] == ["embedding"] + assert len(matrices) == 1 + assert rest == [{"id": i} for i in range(N)] + assert "embedding" in rows[0], "the caller's rows are not changed" + + def test_numpy_rows_and_ints_are_taken_out(self, vec_db): + assert ( + self._split( + vec_db, "Doc", [{"embedding": np.arange(DIM, dtype=np.float64)}] * 3 + ) + is not None + ) + assert ( + self._split( + vec_db, "Doc", [{"embedding": [1, 2, 3]}, {"embedding": [1.5, 2, 3]}] + ) + is not None + ) + + def test_the_inherited_declaration_counts(self, vec_db): + assert self._split(vec_db, "Child", _rows(_doubles())) is not None + + @pytest.mark.parametrize( + "rows", + [ + [{"embedding": [1.0, 2.0]}, {"embedding": [1.0, 2.0, 3.0]}], # ragged + [{"embedding": [1.0, 2.0]}, {"embedding": None}], # a null + [{"embedding": [1.0, 2.0]}, {"id": 3}], # missing in a row + [{"embedding": ["a", "b"]}, {"embedding": ["c", "d"]}], # not numbers + [{"embedding": [True, False]}, {"embedding": [False, True]}], # booleans + [{"embedding": []}, {"embedding": []}], # no dimension + [{"embedding": 1.5}, {"embedding": 2.5}], # a scalar + ], + ) + def test_anything_else_stays_in_the_json(self, vec_db, rows): + assert self._split(vec_db, "Doc", rows) is None + + def test_an_undeclared_property_stays_in_the_json(self, vec_db): + assert self._split(vec_db, "Plain", _rows(_doubles())) is None + assert self._split(vec_db, "NoSuchType", _rows(_doubles())) is None + + +class TestInsertMany: + def test_list_vectors_store_what_the_engine_stores_for_a_list(self, vec_db): + x = _doubles() + assert vec_db.insert_many("Doc", _rows(x)) == N + _assert_float32_of(_stored(vec_db, "Doc"), x) + # the reference: the per-record path, where the engine converts the list itself + with vec_db.transaction(): + for i in range(N): + vec_db.new_document("Parent").set( + "embedding", [float(v) for v in x[i]] + ).set("id", i).save() + reference = { + r.get("id"): np.array(r.get("embedding"), dtype=np.float64) + for r in vec_db.query("sql", "SELECT id, embedding FROM Parent") + } + for i, vec in _stored(vec_db, "Doc").items(): + assert np.array_equal(vec, reference[i]) + + @pytest.mark.parametrize("dtype", [np.float32, np.float64]) + def test_numpy_rows(self, vec_db, dtype): + x = _doubles() + rows = [{"id": i, "embedding": x[i].astype(dtype)} for i in range(N)] + assert vec_db.insert_many("Doc", rows) == N + _assert_float32_of(_stored(vec_db, "Doc"), x) + + def test_an_undeclared_property_keeps_its_doubles(self, vec_db): + x = _doubles() + assert vec_db.insert_many("Plain", _rows(x)) == N + got = _stored(vec_db, "Plain") + for i in range(N): + assert np.array_equal(got[i], x[i]) + + def test_numpy_rows_on_an_undeclared_property(self, vec_db): + x = _doubles() + vec_db.insert_many( + "Plain", + [ + {"id": 0, "embedding": x[0]}, + {"id": 1, "embedding": x[1].astype(np.float32)}, + ], + ) + got = _stored(vec_db, "Plain") + assert np.array_equal( + got[0], x[0] + ) # a float64 array is the list it would have been + assert np.array_equal(got[1], x[1].astype(np.float32).astype(np.float64)) + + def test_ragged_and_null_vectors_keep_the_json_path(self, vec_db): + rows = [ + {"id": 0, "embedding": [1.0, 2.0]}, + {"id": 1, "embedding": [1.0, 2.0, 3.0]}, + {"id": 2, "embedding": None}, + ] + assert vec_db.insert_many("Doc", rows) == 3 + got = { + r["id"]: r["embedding"] + for r in vec_db.query("sql", "SELECT id, embedding FROM Doc").to_list() + } + assert got[0] == [1.0, 2.0] and got[1] == [1.0, 2.0, 3.0] and got[2] is None + + def test_other_properties_still_cross_in_the_json(self, vec_db): + vec_db.command("sql", "CREATE PROPERTY Doc.name STRING") + x = _doubles(5) + vec_db.insert_many( + "Doc", + [{"id": i, "name": f"n{i}", "embedding": x[i].tolist()} for i in range(5)], + ) + names = { + r["id"]: r["name"] + for r in vec_db.query("sql", "SELECT id, name FROM Doc").to_list() + } + assert names == {i: f"n{i}" for i in range(5)} + + def test_a_chunk_boundary_and_a_row_count_that_is_not_a_chunk(self, vec_db): + x = _doubles(23_000, 4) + assert vec_db.insert_many("Doc", _rows(x), commit_every=7_000) == 23_000 + got = _stored(vec_db, "Doc") + assert len(got) == 23_000 + for i in (0, 9_999, 10_000, 22_999): + assert np.array_equal(got[i], x[i].astype(np.float32).astype(np.float64)) + + def test_parallel(self, vec_db): + x = _doubles(200) + assert vec_db.insert_many("Doc", _rows(x), parallel=True) == 200 + _assert_float32_of(_stored(vec_db, "Doc"), x) + + def test_the_vector_survives_a_rolled_back_batch_not_at_all(self, vec_db): + x = _doubles(5) + with pytest.raises(RuntimeError): + with vec_db.transaction(): + vec_db.insert_many("Doc", _rows(x)) + raise RuntimeError("abandon") + assert _stored(vec_db, "Doc") == {} + + +class TestInsertColumns: + def test_a_two_dimensional_float32_column(self, vec_db): + x = _doubles().astype(np.float32) + assert vec_db.insert_columns("Doc", {"id": np.arange(N), "embedding": x}) == N + stored = _stored(vec_db, "Doc") + for i in range(N): + assert np.array_equal(stored[i], x[i].astype(np.float64)) + + def test_parallel(self, vec_db): + x = _doubles(120).astype(np.float32) + assert ( + vec_db.insert_columns( + "Doc", {"id": np.arange(120), "embedding": x}, parallel=True + ) + == 120 + ) + assert len(_stored(vec_db, "Doc")) == 120 + + def test_a_column_of_numpy_rows_is_no_longer_a_type_error(self, vec_db): + x = _doubles(10).astype(np.float32) + vec_db.insert_columns( + "Doc", {"id": np.arange(10), "embedding": [x[i] for i in range(10)]} + ) + assert len(_stored(vec_db, "Doc")) == 10 + + def test_other_two_dimensional_columns_are_refused(self, vec_db): + with pytest.raises(ValueError, match="one-dimensional"): + vec_db.insert_columns( + "Doc", {"id": np.arange(N), "embedding": _doubles()} + ) # float64 + with pytest.raises(ValueError, match="one-dimensional"): + vec_db.insert_columns( + "Doc", + {"id": np.arange(N), "embedding": np.zeros((N, 0), dtype=np.float32)}, + ) + assert ( + vec_db.query("sql", "SELECT count(*) AS c FROM Doc").first().get("c") == 0 + ) + + def test_a_row_count_that_differs_is_refused_before_anything_is_written( + self, vec_db + ): + with pytest.raises(ValueError, match="differ in length"): + vec_db.insert_columns( + "Doc", + {"id": np.arange(N + 1), "embedding": _doubles().astype(np.float32)}, + ) + assert ( + vec_db.query("sql", "SELECT count(*) AS c FROM Doc").first().get("c") == 0 + ) + + +class TestCreateVertices: + def _ids(self, db, type_name="Vx"): + query = f"SELECT id, embedding FROM {type_name}" # nosec B608 - test-owned type name + return { + r["id"]: np.array(r["embedding"], dtype=np.float64) + for r in db.query("sql", query).to_list() + } + + def test_list_vectors(self, vec_db): + x = _doubles() + with vec_db.graph_batch() as gb: + rids = gb.create_vertices("Vx", _rows(x)) + assert len(rids) == N and len(set(rids)) == N + _assert_float32_of(self._ids(vec_db), x) + # the RIDs come back in row order + for i in (0, 17, N - 1): + query = f"SELECT id FROM {rids[i]}" # nosec B608 - a RID from this test + assert vec_db.query("sql", query).first().get("id") == i + + @pytest.mark.parametrize("dtype", [np.float32, np.float64]) + def test_numpy_rows(self, vec_db, dtype): + x = _doubles() + with vec_db.graph_batch() as gb: + gb.create_vertices( + "Vx", [{"id": i, "embedding": x[i].astype(dtype)} for i in range(N)] + ) + _assert_float32_of(self._ids(vec_db), x) + + def test_one_value_json_cannot_carry_sends_only_its_chunk_the_typed_way( + self, vec_db + ): + vec_db.command("sql", "CREATE PROPERTY Vx.born DATETIME") + x = _doubles(10) + rows = _rows(x) + rows[4]["born"] = datetime(2020, 1, 2, 3, 4, 5) + with vec_db.graph_batch() as gb: + rids = gb.create_vertices("Vx", rows) + assert len(rids) == 10 + _assert_float32_of(self._ids(vec_db), x) + assert vec_db.query("sql", "SELECT born FROM Vx WHERE id = 4").first().get( + "born" + ) == datetime(2020, 1, 2, 3, 4, 5) + + def test_rows_with_no_properties_and_none_rows(self, vec_db): + with vec_db.graph_batch() as gb: + assert len(gb.create_vertices("Vx", [None, {}, {"id": 5}])) == 3 + assert vec_db.query("sql", "SELECT count(*) AS c FROM Vx").first().get("c") == 3 + + def test_an_empty_list(self, vec_db): + with vec_db.graph_batch() as gb: + assert gb.create_vertices("Vx", []) == [] + + +class TestDocumentSet: + @pytest.mark.parametrize("dtype", [np.float32, np.float64, np.int64]) + def test_a_numpy_array_is_a_property_value(self, vec_db, dtype): + v = (_doubles(1)[0] * 10).astype(dtype) + with vec_db.transaction(): + vec_db.new_document("Doc").set("id", 1).set("embedding", v).save() + got = _stored(vec_db, "Doc")[1] + assert np.array_equal(got, v.astype(np.float32).astype(np.float64)) + + def test_convert_python_to_java_float32_is_one_float_array(self, vec_db): + v = np.arange(DIM, dtype=np.float32) + java = type_conversion.convert_python_to_java(v) + assert str(java.getClass().getName()) == "[F" + assert list(java) == list(v) + + +class TestAllowedRids: + def _index(self, db): + x = _doubles(40) + with db.graph_batch() as gb: + gb.create_vertices("Vx", _rows(x)) + db.command( + "sql", + "CREATE INDEX ON Vx (embedding) LSM_VECTOR METADATA " + + '{"dimensions": %d, "similarity": "EUCLIDEAN"}' % DIM, + ) + rids = { + r["id"]: str(r["rid"]) + for r in db.query("sql", "SELECT id, @rid AS rid FROM Vx").to_list() + } + return db.schema.get_vector_index("Vx", "embedding"), x, rids + + def test_the_set_holds_one_rid_per_string(self, vec_db): + index, _x, rids = self._index(vec_db) + s = index._build_allowed_rids_set( + list(rids.values())[:7] + [list(rids.values())[0]] + ) + assert int(s.size()) == 7 + assert all(str(r) in set(rids.values()) for r in s) + + def test_text_rids_restrict_the_search(self, vec_db): + index, x, rids = self._index(vec_db) + allowed = [rids[i] for i in (3, 11, 25)] + hits = index.find_nearest(x[11].astype(np.float32), k=5, allowed_rids=allowed) + assert {str(h[0].get_rid()) for h in hits} <= set(allowed) + assert str(hits[0][0].get_rid()) == rids[11] + + def test_rid_objects_and_text_mixed_still_work(self, vec_db): + index, x, rids = self._index(vec_db) + java_rid = vec_db.to_java_rid(rids[11]) + s = index._build_allowed_rids_set([java_rid, rids[3]]) + assert int(s.size()) == 2 diff --git a/tests/test_python_callbacks.py b/tests/test_python_callbacks.py index 53dd205..bb80a4d 100644 --- a/tests/test_python_callbacks.py +++ b/tests/test_python_callbacks.py @@ -183,10 +183,7 @@ def interrupt(value): def test_unconvertible_return_values_are_readable_errors(db): - import numpy as np - for label, value in ( - ("TypeError", np.array([1, 2])), ("OverflowError", 2**63), ("TypeError", object()), ): @@ -200,6 +197,18 @@ def test_unconvertible_return_values_are_readable_errors(db): assert db.query("sql", "SELECT bad(1) AS x").first().get("x") == 2**62 +def test_a_numpy_array_is_a_return_value(db): + """It raised TypeError until numpy arrays became property values (FloatVectors): a float32 vector is one float[], any other dtype the list it would have been.""" + import numpy as np + + db.register_function( + "vec", lambda x: np.array([1.5, 2.5], dtype=np.float32), replace=True + ) + assert db.query("sql", "SELECT vec(1) AS x").first().get("x") == [1.5, 2.5] + db.register_function("ints", lambda x: np.array([1, 2]), replace=True) + assert db.query("sql", "SELECT ints(1) AS x").first().get("x") == [1, 2] + + @pytest.mark.xfail( not getattr(sys, "_is_gil_enabled", lambda: True)(), reason=( diff --git a/tests/test_vector.py b/tests/test_vector.py index 6f6bace..db8ee8e 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -749,7 +749,12 @@ def hits(self, java_db, pairs): assert hit_calls["count"] == 1 def test_lsm_vector_search_uses_database_rid_conversion(self, test_db, monkeypatch): - """RID whitelist conversion should go through the Database wrapper.""" + """RID whitelist conversion should go through the Database wrapper. + + A list of RID strings is the exception: it becomes the whole set in one bridge call + (GraphCalls.ridSet, tested in test_float_vector_bulk.py), so this test passes one RID + as a Java RID object, which makes the list take the per-RID path. + """ test_db.command("sql", "CREATE VERTEX TYPE Doc") test_db.command("sql", "CREATE PROPERTY Doc.id INTEGER") @@ -779,6 +784,7 @@ def test_lsm_vector_search_uses_database_rid_conversion(self, test_db, monkeypat ) ] + mixed_rids = [allowed_rids[0], test_db.to_java_rid(allowed_rids[1])] rid_conversion_calls = [] original_to_java_rid = test_db.to_java_rid @@ -788,7 +794,7 @@ def wrapped_to_java_rid(value): monkeypatch.setattr(test_db, "to_java_rid", wrapped_to_java_rid) - results = index.find_nearest([1.0, 0.0, 0.0], k=2, allowed_rids=allowed_rids) + results = index.find_nearest([1.0, 0.0, 0.0], k=2, allowed_rids=mixed_rids) assert len(results) == 2 assert rid_conversion_calls[: len(allowed_rids)] == allowed_rids