|
| 1 | +"""insert_arrow and insert_dataframe with a vector column on the JVM wheel: the column crosses as one float32 matrix, not as one Python list of |
| 2 | +Python floats per row, and what is stored is unchanged. |
| 3 | +
|
| 4 | +A ``FixedSizeList<float32>`` (or list, or large list) column into an ``ARRAY_OF_FLOATS`` property went ``chunk.to_pylist()`` (one Python float |
| 5 | +per element, 31.5 MiB of Python heap for one 10,000 x 96 batch against a 3.7 MiB flat buffer), then ``_column_to_java`` converted every row |
| 6 | +with ``convert_python_to_java`` (an ``ArrayList<Double>`` of 96 boxed doubles per row, 100,000 conversions for the 10 x 10,000 rows of the |
| 7 | +study) and the engine narrowed each list back to a ``float[]``. In the untimed profile that conversion was about 63 percent of the call, |
| 8 | +and the boxed doubles are what the JVM heap held (``heap used`` 183 to 195 MiB at the end against 69 MiB live after a collection). Now a |
| 9 | +float32 column with one length and no null stays an Arrow buffer, ``_float32_matrix`` lays it out as one 2-D float32 numpy matrix, and |
| 10 | +``_column_to_java`` sends it as one ``float[]`` inside the bridge's ``FloatVectors.Matrix`` (the route a 2-D numpy column of |
| 11 | +``insert_columns`` already took); the documents get a row each Java-side. |
| 12 | +
|
| 13 | +Three tests fail on the code before (they count the per-row conversions, the matrix crossings and the Python heap of one batch); the others |
| 14 | +pin what the old route stored, so the new one must store the same: count, an order-independent digest of every row's vector bytes read back |
| 15 | +through SQL and through the result API, both again after close and reopen, the special float32 values, a null row, a null element, vectors |
| 16 | +of different length, a wrong dimension against a vector index, float64 and integer lists (which are NOT taken: the old route refuses a |
| 17 | +value past the float range and rounds the rest, a float32 cast would store an infinity), and the properties that are not declared |
| 18 | +``ARRAY_OF_FLOATS`` (they keep a list). The JVM build only: it counts JPype crossings and reads Java-side values, so conftest.py does not |
| 19 | +collect this file on the native wheel (whose vector columns stay Arrow arrays, test_native_columnar_memory.py). |
| 20 | +""" |
| 21 | + |
| 22 | +import hashlib |
| 23 | +import tracemalloc |
| 24 | + |
| 25 | +import numpy as np |
| 26 | +import pyarrow as pa |
| 27 | +import pytest |
| 28 | + |
| 29 | +import arcadedb |
| 30 | +from arcadedb import core |
| 31 | +from arcadedb.exceptions import ArcadeDBError |
| 32 | + |
| 33 | +F32 = np.float32 |
| 34 | +SHAPES = ["fixed_size", "list", "large_list"] |
| 35 | +TYPES = ["VERTEX", "DOCUMENT"] |
| 36 | + |
| 37 | + |
| 38 | +def _matrix(n, dim, start=0, specials=True): |
| 39 | + """float32 vectors, a pure function of the row number, with the values a decimal text of 6 digits would not hold in rows 0 and 1.""" |
| 40 | + r = np.arange(start, start + n, dtype=np.int64) |
| 41 | + j = np.arange(dim, dtype=np.int64) |
| 42 | + m = (((r[:, None] * 257 + j[None, :] * 17) % 1000).astype(F32)) / F32(1000.0) |
| 43 | + if specials and start == 0 and n > 1 and dim > 3: |
| 44 | + m[0, 0] = F32(-0.0) |
| 45 | + m[1, 0] = np.finfo(F32).max |
| 46 | + m[1, 1] = np.finfo(F32).tiny |
| 47 | + m[1, 2] = np.finfo(F32).smallest_subnormal |
| 48 | + return m |
| 49 | + |
| 50 | + |
| 51 | +def _column(matrix, shape): |
| 52 | + dim = matrix.shape[1] |
| 53 | + flat = pa.array(matrix.reshape(-1), type=pa.float32()) |
| 54 | + if shape == "fixed_size": |
| 55 | + return pa.FixedSizeListArray.from_arrays(flat, dim) |
| 56 | + kind = pa.list_ if shape == "list" else pa.large_list |
| 57 | + rows = [row.tolist() for row in matrix] |
| 58 | + return pa.array(rows, type=kind(pa.float32())) |
| 59 | + |
| 60 | + |
| 61 | +def _table(matrix, shape, start=0): |
| 62 | + n = matrix.shape[0] |
| 63 | + return pa.table({"k": pa.array(np.arange(start, start + n), type=pa.int32()), "embedding": _column(matrix, shape)}) |
| 64 | + |
| 65 | + |
| 66 | +def _make(db, name, kind="VERTEX", embedding="ARRAY_OF_FLOATS", index=None): |
| 67 | + db.command("sql", f"CREATE {kind} TYPE {name}") |
| 68 | + db.command("sql", f"CREATE PROPERTY {name}.k INTEGER") |
| 69 | + if embedding: |
| 70 | + db.command("sql", f"CREATE PROPERTY {name}.embedding {embedding}") |
| 71 | + if index: |
| 72 | + db.command("sql", f"CREATE INDEX ON {name} (embedding) LSM_VECTOR METADATA {{dimensions: {index}, similarity: 'COSINE'}}") |
| 73 | + |
| 74 | + |
| 75 | +def _count(db, name): |
| 76 | + return db.query("sql", f"SELECT count(*) AS c FROM {name}").first().get("c") # nosec B608 - test-owned name |
| 77 | + |
| 78 | + |
| 79 | +def _key(k, vector_bytes): |
| 80 | + return int.from_bytes(hashlib.sha256(int(k).to_bytes(8, "little") + vector_bytes).digest()[:16], "little") |
| 81 | + |
| 82 | + |
| 83 | +def _expected_digest(matrix, start=0): |
| 84 | + """The digest the stored table must have, from the generator, never from the database.""" |
| 85 | + h = 0 |
| 86 | + for i, row in enumerate(matrix): |
| 87 | + h ^= _key(start + i, row.tobytes()) |
| 88 | + return h |
| 89 | + |
| 90 | + |
| 91 | +def _sql_digest(db, name): |
| 92 | + h, n = 0, 0 |
| 93 | + for row in db.query("sql", f"SELECT k, embedding FROM {name}").to_list(): # nosec B608 - test-owned name |
| 94 | + h ^= _key(row["k"], np.asarray(row["embedding"], dtype=np.float64).astype(F32).tobytes()) |
| 95 | + n += 1 |
| 96 | + return n, h |
| 97 | + |
| 98 | + |
| 99 | +def _api_digest(db, name): |
| 100 | + """The same digest through the result API (one Result per record, ``get`` per property), not through ``to_list``.""" |
| 101 | + h, n = 0, 0 |
| 102 | + result_set = db.query("sql", f"SELECT FROM {name}") # nosec B608 - test-owned name |
| 103 | + for result in result_set: |
| 104 | + h ^= _key(result.get("k"), np.asarray(result.get("embedding"), dtype=np.float64).astype(F32).tobytes()) |
| 105 | + n += 1 |
| 106 | + return n, h |
| 107 | + |
| 108 | + |
| 109 | +def _bits(values): |
| 110 | + return np.asarray(values, dtype=np.float64).astype(F32).view(np.uint32).tolist() |
| 111 | + |
| 112 | + |
| 113 | +@pytest.fixture |
| 114 | +def crossings(monkeypatch): |
| 115 | + """What one load does on the JPype side: the per-value conversions and the bridge's Matrix constructions.""" |
| 116 | + log = {"convert": [], "matrix": []} |
| 117 | + real_convert = core.convert_python_to_java |
| 118 | + real_class = core._java_class |
| 119 | + |
| 120 | + def convert(value): |
| 121 | + log["convert"].append(type(value).__name__) |
| 122 | + return real_convert(value) |
| 123 | + |
| 124 | + class _Counting: |
| 125 | + def __init__(self, target): |
| 126 | + self._target = target |
| 127 | + |
| 128 | + def __call__(self, *args): |
| 129 | + log["matrix"].append(len(args[0])) |
| 130 | + return self._target(*args) |
| 131 | + |
| 132 | + def __getattr__(self, name): |
| 133 | + return getattr(self._target, name) |
| 134 | + |
| 135 | + def java_class(name): |
| 136 | + cls = real_class(name) |
| 137 | + return _Counting(cls) if name == "com.arcadedb.python.FloatVectors$Matrix" else cls |
| 138 | + |
| 139 | + monkeypatch.setattr(core, "convert_python_to_java", convert) |
| 140 | + monkeypatch.setattr(core, "_java_class", java_class) |
| 141 | + return log |
| 142 | + |
| 143 | + |
| 144 | +@pytest.mark.parametrize("kind", TYPES) |
| 145 | +@pytest.mark.parametrize("shape", SHAPES) |
| 146 | +def test_count_digest_and_reopen_of_a_batched_vector_load(temp_db, temp_db_path_of, shape, kind): |
| 147 | + """1,300 rows in batches of 500 (500, 500, 300: the batches are slices of one table): count, SQL digest, API digest, then the same after reopen.""" |
| 148 | + n, dim = 1_300, 24 |
| 149 | + matrix = _matrix(n, dim) |
| 150 | + _make(temp_db, "Vec", kind) |
| 151 | + assert temp_db.insert_arrow("Vec", _table(matrix, shape), batch_size=500) == n |
| 152 | + want = (n, _expected_digest(matrix)) |
| 153 | + assert _count(temp_db, "Vec") == n |
| 154 | + assert _sql_digest(temp_db, "Vec") == want |
| 155 | + assert _api_digest(temp_db, "Vec") == want |
| 156 | + path = temp_db_path_of(temp_db) |
| 157 | + temp_db.close() |
| 158 | + reopened = arcadedb.open_database(path) |
| 159 | + try: |
| 160 | + assert _count(reopened, "Vec") == n |
| 161 | + assert _sql_digest(reopened, "Vec") == want |
| 162 | + assert _api_digest(reopened, "Vec") == want |
| 163 | + finally: |
| 164 | + reopened.close() |
| 165 | + |
| 166 | + |
| 167 | +@pytest.fixture |
| 168 | +def temp_db_path_of(): |
| 169 | + def path_of(db): |
| 170 | + return db.get_java_database().getDatabasePath() |
| 171 | + |
| 172 | + return path_of |
| 173 | + |
| 174 | + |
| 175 | +def test_a_load_in_several_chunks_of_one_table_keeps_every_row(temp_db): |
| 176 | + """Two tables concatenated (a ChunkedArray of two chunks) cut into batches of 4 (offsets that start inside a chunk).""" |
| 177 | + a, b = _matrix(3, 6, 0), _matrix(4, 6, 3) |
| 178 | + _make(temp_db, "Vec") |
| 179 | + table = pa.concat_tables([_table(a, "fixed_size", 0), _table(b, "fixed_size", 3)]) |
| 180 | + assert temp_db.insert_arrow("Vec", table, batch_size=4) == 7 |
| 181 | + assert _sql_digest(temp_db, "Vec") == (7, _expected_digest(a, 0) ^ _expected_digest(b, 3)) |
| 182 | + |
| 183 | + |
| 184 | +@pytest.mark.parametrize("shape", SHAPES) |
| 185 | +def test_the_special_float32_values_are_stored_bit_for_bit(temp_db, shape): |
| 186 | + special = np.array([[-0.0, np.finfo(F32).max, np.finfo(F32).tiny, np.finfo(F32).smallest_subnormal, np.nan, np.inf, -np.inf, 16777217.0]], dtype=F32) |
| 187 | + _make(temp_db, "Vec") |
| 188 | + assert temp_db.insert_arrow("Vec", _table(special, shape)) == 1 |
| 189 | + stored = temp_db.query("sql", "SELECT embedding FROM Vec").first().get("embedding") |
| 190 | + assert len(stored) == 8 |
| 191 | + got = np.asarray(stored, dtype=np.float64).astype(F32) |
| 192 | + assert np.isnan(got[4]) |
| 193 | + keep = [0, 1, 2, 3, 5, 6, 7] |
| 194 | + assert got.view(np.uint32)[keep].tolist() == special.view(np.uint32)[0, keep].tolist() |
| 195 | + assert got[7] == F32(16777216.0) # 16777217.0 is not a float32; the row holds the float32 that was written |
| 196 | + |
| 197 | + |
| 198 | +def test_a_null_row_stays_null_and_its_neighbours_are_exact(temp_db): |
| 199 | + _make(temp_db, "Vec") |
| 200 | + table = pa.table({"k": pa.array([0, 1, 2], pa.int32()), "embedding": pa.array([[1.0, 2.0], None, [3.0, 4.0]], pa.list_(pa.float32(), 2))}) |
| 201 | + assert temp_db.insert_arrow("Vec", table) == 3 |
| 202 | + rows = {r["k"]: r["embedding"] for r in temp_db.query("sql", "SELECT k, embedding FROM Vec").to_list()} |
| 203 | + assert rows[1] is None |
| 204 | + assert _bits(rows[0]) == _bits([1.0, 2.0]) and _bits(rows[2]) == _bits([3.0, 4.0]) |
| 205 | + |
| 206 | + |
| 207 | +def test_a_null_element_is_refused_and_the_batch_is_rolled_back(temp_db): |
| 208 | + _make(temp_db, "Vec") |
| 209 | + table = pa.table({"k": pa.array([0, 1], pa.int32()), "embedding": pa.array([[1.0, None, 3.0], [4.0, 5.0, 6.0]], pa.list_(pa.float32()))}) |
| 210 | + with pytest.raises(ArcadeDBError, match="null element"): |
| 211 | + temp_db.insert_arrow("Vec", table) |
| 212 | + assert _count(temp_db, "Vec") == 0 |
| 213 | + |
| 214 | + |
| 215 | +def test_vectors_of_different_length_are_stored_as_given_without_an_index_and_refused_with_one(temp_db): |
| 216 | + _make(temp_db, "Free") |
| 217 | + ragged = pa.table({"k": pa.array([0, 1, 2], pa.int32()), "embedding": pa.array([[1.0, 2.0], [1.0, 2.0, 3.0], [4.0]], pa.list_(pa.float32()))}) |
| 218 | + assert temp_db.insert_arrow("Free", ragged) == 3 |
| 219 | + rows = {r["k"]: r["embedding"] for r in temp_db.query("sql", "SELECT k, embedding FROM Free").to_list()} |
| 220 | + assert {k: len(v) for k, v in rows.items()} == {0: 2, 1: 3, 2: 1} |
| 221 | + _make(temp_db, "Indexed", index=2) |
| 222 | + with pytest.raises(ArcadeDBError, match="Vector dimension does not match index dimension 2: got float.. of length 3"): |
| 223 | + temp_db.insert_arrow("Indexed", ragged.slice(0, 2)) |
| 224 | + assert _count(temp_db, "Indexed") == 0 |
| 225 | + |
| 226 | + |
| 227 | +def test_a_fixed_size_vector_of_the_wrong_dimension_is_refused_by_the_vector_index(temp_db): |
| 228 | + _make(temp_db, "Indexed", index=2) |
| 229 | + wrong = _table(_matrix(2, 3, specials=False), "fixed_size") |
| 230 | + with pytest.raises(ArcadeDBError, match="Vector dimension does not match index dimension 2: got float.. of length 3"): |
| 231 | + temp_db.insert_arrow("Indexed", wrong) |
| 232 | + assert _count(temp_db, "Indexed") == 0 |
| 233 | + right = _matrix(2, 2, specials=False) |
| 234 | + assert temp_db.insert_arrow("Indexed", _table(right, "fixed_size")) == 2 |
| 235 | + assert _sql_digest(temp_db, "Indexed") == (2, _expected_digest(right)) |
| 236 | + |
| 237 | + |
| 238 | +def test_float64_and_integer_lists_keep_what_the_old_route_stored(temp_db): |
| 239 | + """Not taken by the matrix route: a float64 value past the float range is refused (a cast would store an infinity), the rest rounds.""" |
| 240 | + _make(temp_db, "Dbl") |
| 241 | + ok = pa.table({"k": pa.array([0], pa.int32()), "embedding": pa.array([[0.1, 16777217.0, -2.5]], pa.list_(pa.float64()))}) |
| 242 | + assert temp_db.insert_arrow("Dbl", ok) == 1 |
| 243 | + got = temp_db.query("sql", "SELECT embedding FROM Dbl").first().get("embedding") |
| 244 | + assert _bits(got) == _bits([0.1, 16777217.0, -2.5]) |
| 245 | + assert got[1] == 16777216.0 |
| 246 | + _make(temp_db, "Over") |
| 247 | + over = pa.table({"k": pa.array([0], pa.int32()), "embedding": pa.array([[1e39, 1.0]], pa.list_(pa.float64()))}) |
| 248 | + with pytest.raises(ArcadeDBError, match="1.0E39"): |
| 249 | + temp_db.insert_arrow("Over", over) |
| 250 | + assert _count(temp_db, "Over") == 0 |
| 251 | + _make(temp_db, "Ints") |
| 252 | + ints = pa.table({"k": pa.array([0], pa.int32()), "embedding": pa.array([[16777217, 2**40, -3]], pa.list_(pa.int64()))}) |
| 253 | + assert temp_db.insert_arrow("Ints", ints) == 1 |
| 254 | + assert temp_db.query("sql", "SELECT embedding FROM Ints").first().get("embedding") == [16777216.0, 1099511627776.0, -3.0] |
| 255 | + |
| 256 | + |
| 257 | +@pytest.mark.parametrize("embedding", [None, "LIST"], ids=["undeclared", "declared_list"]) |
| 258 | +def test_a_property_that_is_not_array_of_floats_keeps_a_list_and_takes_no_matrix(temp_db, crossings, embedding): |
| 259 | + _make(temp_db, "Free", embedding=embedding) |
| 260 | + matrix = _matrix(5, 4, specials=False) |
| 261 | + assert temp_db.insert_arrow("Free", _table(matrix, "fixed_size")) == 5 |
| 262 | + assert crossings["matrix"] == [] |
| 263 | + assert _sql_digest(temp_db, "Free") == (5, _expected_digest(matrix)) |
| 264 | + |
| 265 | + |
| 266 | +def test_empty_vectors_are_stored_as_empty_lists(temp_db): |
| 267 | + _make(temp_db, "Vec") |
| 268 | + table = pa.table({"k": pa.array([0, 1], pa.int32()), "embedding": pa.array([[], []], pa.list_(pa.float32()))}) |
| 269 | + assert temp_db.insert_arrow("Vec", table) == 2 |
| 270 | + assert [r["embedding"] for r in temp_db.query("sql", "SELECT embedding FROM Vec ORDER BY k").to_list()] == [[], []] |
| 271 | + |
| 272 | + |
| 273 | +# ---- the three tests that fail on the code before the fix: they count what the load does, never how long it takes |
| 274 | + |
| 275 | + |
| 276 | +@pytest.mark.parametrize("shape", SHAPES) |
| 277 | +def test_a_vector_column_converts_no_python_value_per_row(temp_db, crossings, shape): |
| 278 | + """Before: one ``convert_python_to_java`` per row (an ArrayList of boxed doubles). After: the vector column makes none.""" |
| 279 | + n = 600 |
| 280 | + _make(temp_db, "Vec") |
| 281 | + assert temp_db.insert_arrow("Vec", _table(_matrix(n, 8), shape), batch_size=n) == n |
| 282 | + assert crossings["convert"] == [] |
| 283 | + |
| 284 | + |
| 285 | +def test_a_vector_column_crosses_as_one_matrix_per_batch(temp_db, crossings): |
| 286 | + """1,300 rows in batches of 500: three batches, so three ``FloatVectors.Matrix`` of one flat float[] each (500 x 24, 500 x 24, 300 x 24).""" |
| 287 | + _make(temp_db, "Vec") |
| 288 | + assert temp_db.insert_arrow("Vec", _table(_matrix(1_300, 24), "fixed_size"), batch_size=500) == 1_300 |
| 289 | + assert crossings["matrix"] == [500 * 24, 500 * 24, 300 * 24] |
| 290 | + |
| 291 | + |
| 292 | +def test_the_python_heap_of_one_batch_stays_below_twice_the_flat_buffer(temp_db): |
| 293 | + """10,000 x 96 float32 is a 3.66 MiB buffer; the per-row Python lists of floats were 31.5 MiB at their peak (tracemalloc sees Python and |
| 294 | + numpy allocations, not the Java heap and not Arrow's pool). The bound is the buffer twice, nothing about time.""" |
| 295 | + n, dim = 10_000, 96 |
| 296 | + table = _table(_matrix(n, dim, specials=False), "fixed_size") |
| 297 | + _make(temp_db, "Vec") |
| 298 | + _make(temp_db, "Twin") |
| 299 | + assert temp_db.insert_arrow("Twin", table.slice(0, 500)) == 500 # first-call imports and caches stay out of the measurement |
| 300 | + flat = n * dim * 4 |
| 301 | + tracemalloc.start() |
| 302 | + try: |
| 303 | + tracemalloc.reset_peak() |
| 304 | + base = tracemalloc.get_traced_memory()[0] |
| 305 | + assert temp_db.insert_arrow("Vec", table, batch_size=n) == n |
| 306 | + peak = tracemalloc.get_traced_memory()[1] - base |
| 307 | + finally: |
| 308 | + tracemalloc.stop() |
| 309 | + assert peak < 2 * flat, f"{peak} bytes of Python heap for a {flat} byte vector column" |
0 commit comments