Skip to content

Commit 56569a0

Browse files
committed
integration: merge perf/jvm-arrow-vector-conversion (9914e93: float32 vector columns of an Arrow vertex load cross to Java as one matrix per batch instead of one boxed list per row, JVM wheel only; timed A/B arrow-vector-fix-ab-20261110: the l3d WAL-off cell is a WIN at B/A 0.154, the other gating cells are not decided by the pre-registered rule; refs #347 #335)
2 parents ece871a + 9914e93 commit 56569a0

3 files changed

Lines changed: 363 additions & 1 deletion

File tree

‎src/arcadedb/dataframe.py‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -336,10 +336,57 @@ def _is_number_list(pa, arrow_type) -> bool:
336336
)
337337

338338

339-
def _column_values(name: str, chunk, pa, native: bool = False, arrow: bool = False):
339+
_JAVA_INT_MAX = 2**31 - 1
340+
341+
342+
def _float32_matrix(pa, chunk):
343+
"""A list, large list or fixed-size list column of float32 with no null row, no null element and one length, as one 2-D float32
344+
numpy matrix (a view of Arrow's value buffer: ``flatten`` follows a slice's offset), else None. None means "not this shape" and
345+
leaves the column to ``to_pylist``, which keeps what that route stored or refused: a null row stays null, a null element is
346+
refused, vectors of different length are stored as given or refused by a vector index. Only float32 is taken: a float64 or integer
347+
value is narrowed by the engine (and refused past the float range), which a cast here would not do.
348+
"""
349+
types = pa.types
350+
arrow_type = chunk.type
351+
if not (
352+
types.is_list(arrow_type)
353+
or types.is_large_list(arrow_type)
354+
or types.is_fixed_size_list(arrow_type)
355+
):
356+
return None
357+
n = len(chunk)
358+
if not types.is_float32(arrow_type.value_type) or not n or chunk.null_count:
359+
return None
360+
flat = chunk.flatten()
361+
if flat.null_count or len(flat) % n or len(flat) > _JAVA_INT_MAX:
362+
return None
363+
dim = len(flat) // n
364+
if dim == 0:
365+
return None
366+
if not types.is_fixed_size_list(arrow_type):
367+
import pyarrow.compute as pc
368+
369+
lengths = pc.list_value_length(chunk)
370+
if pc.min(lengths).as_py() != dim or pc.max(lengths).as_py() != dim:
371+
return None
372+
return flat.to_numpy(zero_copy_only=False).reshape(n, dim)
373+
374+
375+
def _column_values(
376+
name: str,
377+
chunk,
378+
pa,
379+
native: bool = False,
380+
arrow: bool = False,
381+
float_vectors: bool = False,
382+
):
340383
"""One Arrow array as what ``Database.insert_columns`` takes: a numpy array when the column crosses as one
341384
buffer (integers, floats and booleans without nulls), else a list of values (``None`` is a null).
342385
386+
``float_vectors`` (the JVM backend, a property declared ``ARRAY_OF_FLOATS``): a float32 vector column of one length and no null is
387+
one 2-D float32 matrix, which ``insert_columns`` sends as one ``float[]``; ``to_pylist`` made one Python float per element and
388+
``insert_columns`` one ``ArrayList`` of boxed doubles per row (see ``_float32_matrix`` for what is left to it).
389+
343390
``arrow`` (the native backend's document and vertex load): a string, binary, timestamp, date or decimal column, a list of
344391
numbers (a vector) and a numeric column with nulls, stay Arrow arrays, which the columnar entry reads in place (the offsets,
345392
the bytes, the integer ticks, the vector values and the validity bitmap) instead of as Python values; ``_insert_documents_native``
@@ -382,6 +429,10 @@ def _column_values(name: str, chunk, pa, native: bool = False, arrow: bool = Fal
382429
return chunk
383430
if arrow and kind == "list" and _is_number_list(pa, chunk.type):
384431
return chunk # a vector column: the columnar entry reads the values in place (to_pylist makes one Python float per element)
432+
if float_vectors and kind == "list":
433+
matrix = _float32_matrix(pa, chunk)
434+
if matrix is not None:
435+
return matrix
385436
if kind in ("timestamp", "date"):
386437
return (
387438
_temporal_to_tagged(chunk, pa) if native else _temporal_to_java(chunk, pa)
@@ -462,6 +513,7 @@ def insert_table(
462513
pa,
463514
native,
464515
arrow=native,
516+
float_vectors=not native and declared.get(name) == "ARRAY_OF_FLOATS",
465517
)
466518
for i, name in enumerate(names)
467519
}

‎tests/conftest.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@
149149
"test_graph_batch_single_pass.py": "the JVM wheel's graph batch bulk path (GraphBatch._create_vertices_json_bulk, the property-matrix fallback)",
150150
"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)",
151151
"test_insert_columns_string_blob.py": "the bridge's StringColumn and the JVM wheel's core._column_to_java (JPype arrays, a spy on the per-element conversion)",
152+
"test_arrow_vector_column_jvm.py": "the JPype crossings of the JVM wheel's Arrow vector column (core.convert_python_to_java, the bridge's FloatVectors.Matrix) and its Python heap",
152153
"test_async_proxy_cache.py": "the JPype proxy classes of the JVM build's async executor",
153154
"test_read_crossings.py": "crossing counts of the JVM build's read path (JPype proxies of the Java objects and the bridge classes)",
154155
"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)",
Lines changed: 309 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,309 @@
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

Comments
 (0)