Skip to content

Commit dbb826f

Browse files
tae898claude
andcommitted
integration: merge perf/319-vertex-loads-columnar again (#335 447a950: Arrow vector columns stay Arrow on native, batches dropped after use)
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2 parents e773cb0 + 447a950 commit dbb826f

3 files changed

Lines changed: 326 additions & 4 deletions

File tree

‎src/arcadedb/dataframe.py‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -322,14 +322,28 @@ def _scalars_to_tagged(kind: str, chunk):
322322
]
323323

324324

325+
def _is_number_list(pa, arrow_type) -> bool:
326+
"""A list, large list or fixed-size list of integers or floats: the column the columnar entry lays out as a float32 matrix."""
327+
types = pa.types
328+
if not (
329+
types.is_list(arrow_type)
330+
or types.is_large_list(arrow_type)
331+
or types.is_fixed_size_list(arrow_type)
332+
):
333+
return False
334+
return types.is_floating(arrow_type.value_type) or types.is_integer(
335+
arrow_type.value_type
336+
)
337+
338+
325339
def _column_values(name: str, chunk, pa, native: bool = False, arrow: bool = False):
326340
"""One Arrow array as what ``Database.insert_columns`` takes: a numpy array when the column crosses as one
327341
buffer (integers, floats and booleans without nulls), else a list of values (``None`` is a null).
328342
329-
``arrow`` (the native backend's document load): a string, binary, timestamp, date or decimal column, and a numeric one with
330-
nulls, stay Arrow arrays, which the columnar entry reads in place (the offsets, the bytes, the integer ticks and the validity
331-
bitmap) instead of as Python values; ``_insert_documents_native`` turns them into tagged values for the row path when a batch
332-
cannot go columnar.
343+
``arrow`` (the native backend's document and vertex load): a string, binary, timestamp, date or decimal column, a list of
344+
numbers (a vector) and a numeric column with nulls, stay Arrow arrays, which the columnar entry reads in place (the offsets,
345+
the bytes, the integer ticks, the vector values and the validity bitmap) instead of as Python values; ``_insert_documents_native``
346+
turns them into Python or tagged values for the row path when a batch cannot go columnar.
333347
"""
334348
types = pa.types
335349
if types.is_dictionary(chunk.type): # a pandas categorical or a Polars Categorical
@@ -366,6 +380,8 @@ def _column_values(name: str, chunk, pa, native: bool = False, arrow: bool = Fal
366380
return chunk.to_pylist()
367381
if arrow and kind in ("timestamp", "date", "decimal", "binary"):
368382
return chunk
383+
if arrow and kind == "list" and _is_number_list(pa, chunk.type):
384+
return chunk # a vector column: the columnar entry reads the values in place (to_pylist makes one Python float per element)
369385
if kind in ("timestamp", "date"):
370386
return (
371387
_temporal_to_tagged(chunk, pa) if native else _temporal_to_java(chunk, pa)
@@ -480,6 +496,10 @@ def insert_table(
480496
) from exc
481497
raise
482498
written += batch.num_rows
499+
del (
500+
columns,
501+
batch,
502+
) # sent: drop this batch's columns before the next batch is built, so two batches are never alive together
483503
return written
484504

485505

‎tests/conftest.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@
5959
collect_ignore.append("test_native_columnar_temporal_review_fixes.py")
6060
# a vertex load by insert_dataframe or insert_arrow crosses as columns, one crossing per batch, observed on the native library
6161
collect_ignore.append("test_native_vertex_columnar.py")
62+
# a document or vertex load keeps its transient Python memory to one batch and a vector column is never turned into Python floats
63+
collect_ignore.append("test_native_columnar_memory.py")
6264
# the rows of an insert_many chunk with date-times, Decimals, bytes and numpy scalars are told apart by the entry points they cross
6365
collect_ignore.append("test_native_mixed_chunk.py")
6466
# the native build's ingest defaults (insert_many and graph_batch commit size), observed on the native library
Lines changed: 300 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,300 @@
1+
"""Native: a document or vertex load by ``insert_arrow`` or ``insert_dataframe`` keeps its transient memory to one batch, and a vector column
2+
is never turned into Python floats on its way to the columnar entry (ArcadeData/arcadedb-python #333 and #335, the extra anonymous memory of
3+
the post1 candidate's native wheel).
4+
5+
Three properties are pinned, none of them with an absolute size:
6+
7+
* ``_column_values(..., arrow=True)`` returns an Arrow list of numbers (a vector column) as the Arrow array, which ``columnar.pack`` reads in
8+
place; the JVM build (``arrow=False``) still gets Python lists. One ``to_pylist()`` of a 64-dimensional vector column makes 64 Python
9+
floats per row (about 2 KB a row), which a 10,000-row batch held for the whole crossing.
10+
* the Python-side peak of a whole load is a small fraction of what one batch of Python floats costs (``tracemalloc`` sees numpy and Python
11+
objects, not Arrow's memory or the image's heap, which is why the ratio is generous).
12+
* ``columnar.pack`` is called once per batch and never with more than ``batch_size`` rows, and what it laid out is released when the batch
13+
is sent (the arrays of batch k are gone before batch k + 1 is packed).
14+
15+
Needs the native library, so conftest.py collects this file on the native backend only.
16+
"""
17+
18+
import ctypes
19+
import gc
20+
import tracemalloc
21+
import weakref
22+
23+
import numpy as np
24+
import pyarrow as pa
25+
import pytest
26+
27+
from arcadedb import dataframe
28+
from arcadedb._native import columnar
29+
30+
DIM = 64
31+
32+
33+
def _vectors(n, start=0, dim=DIM):
34+
r = np.arange(start, start + n, dtype=np.int64)
35+
j = np.arange(dim, dtype=np.int64)
36+
return (
37+
((r[:, None] * 131 + j[None, :] * 17) % 1000).astype(np.float32)
38+
) / np.float32(1000.0)
39+
40+
41+
def _fixed(mat):
42+
return pa.FixedSizeListArray.from_arrays(
43+
pa.array(mat.reshape(-1), type=pa.float32()), mat.shape[1]
44+
)
45+
46+
47+
def _bytes(addr_len):
48+
addr, length = addr_len
49+
return ctypes.string_at(addr, length) if length else b""
50+
51+
52+
def _peak(fn):
53+
gc.collect()
54+
tracemalloc.start()
55+
try:
56+
tracemalloc.reset_peak()
57+
base = tracemalloc.get_traced_memory()[0]
58+
fn()
59+
return tracemalloc.get_traced_memory()[1] - base
60+
finally:
61+
tracemalloc.stop()
62+
63+
64+
@pytest.mark.parametrize(
65+
"make",
66+
[
67+
lambda m: _fixed(m),
68+
lambda m: pa.array(m.astype(np.float64).tolist(), type=pa.list_(pa.float64())),
69+
lambda m: pa.array(m.astype(np.int32).tolist(), type=pa.large_list(pa.int32())),
70+
],
71+
ids=["fixed_size_list_f32", "list_f64", "large_list_i32"],
72+
)
73+
def test_a_vector_column_stays_an_arrow_array_on_the_native_load(make):
74+
chunk = make(_vectors(10))
75+
kept = dataframe._column_values("embedding", chunk, pa, native=True, arrow=True)
76+
assert isinstance(kept, pa.Array) and kept.type == chunk.type
77+
assert (
78+
dataframe._column_values("embedding", chunk, pa, native=True, arrow=False)
79+
== chunk.to_pylist()
80+
)
81+
assert isinstance(
82+
dataframe._column_values("embedding", chunk, pa, native=False, arrow=False),
83+
list,
84+
)
85+
86+
87+
def test_a_list_that_is_not_a_list_of_numbers_still_becomes_python_values():
88+
for chunk in (
89+
pa.array([["a", "b"], ["c"]]),
90+
pa.array([[{"x": 1}], [{"x": 2}]]),
91+
pa.array([[[1.0], [2.0]], [[3.0]]]),
92+
):
93+
got = dataframe._column_values("c", chunk, pa, native=True, arrow=True)
94+
assert isinstance(got, list) and got == chunk.to_pylist()
95+
96+
97+
def test_the_vector_layout_from_arrow_is_the_layout_from_python_lists():
98+
mat = _vectors(300)
99+
table = pa.table({"embedding": _fixed(mat)}).slice(
100+
7, 211
101+
) # a slice: the offset of the array is not zero
102+
declared = columnar.Declared(known={"embedding": "ARRAY_OF_FLOATS"})
103+
kept = dataframe._column_values(
104+
"embedding", table.column(0).combine_chunks(), pa, True, True
105+
)
106+
listed = dataframe._column_values(
107+
"embedding", table.column(0).combine_chunks(), pa, True, False
108+
)
109+
rows_a, (col_a,) = columnar.pack({"embedding": kept}, declared)
110+
rows_b, (col_b,) = columnar.pack({"embedding": listed}, declared)
111+
assert (
112+
rows_a == rows_b == 211
113+
and col_a.kind == col_b.kind == "f"
114+
and col_a.dim == col_b.dim == DIM
115+
)
116+
assert _bytes(col_a.data) == _bytes(col_b.data) == mat[7:218].tobytes()
117+
118+
119+
def test_the_python_side_peak_of_a_vector_batch_is_a_fraction_of_the_python_float_copy():
120+
batch = _fixed(_vectors(10_000))
121+
declared = columnar.Declared(known={"embedding": "ARRAY_OF_FLOATS"})
122+
as_python = _peak(
123+
lambda: columnar.pack(
124+
{
125+
"embedding": dataframe._column_values(
126+
"embedding", batch, pa, True, False
127+
)
128+
},
129+
declared,
130+
)
131+
)
132+
in_place = _peak(
133+
lambda: columnar.pack(
134+
{"embedding": dataframe._column_values("embedding", batch, pa, True, True)},
135+
declared,
136+
)
137+
)
138+
assert (
139+
as_python > 10 * 1024 * 1024
140+
), as_python # 640,000 Python floats: the copy this change removes (about 20 MiB)
141+
assert in_place * 8 < as_python, (
142+
in_place,
143+
as_python,
144+
) # measured: about 100 times less
145+
146+
147+
class _PackSpy:
148+
"""Wraps ``columnar.pack``: the row count of every call and a weak reference to every array it laid out."""
149+
150+
def __init__(self, monkeypatch):
151+
self.rows, self.arrays, self.alive_at_next_call = [], [], []
152+
real = columnar.pack
153+
154+
def pack(columns, declared):
155+
gc.collect()
156+
self.alive_at_next_call.append(
157+
sum(ref() is not None for refs in self.arrays for ref in refs)
158+
)
159+
out = real(columns, declared)
160+
self.rows.append(None if out is None else out[0])
161+
self.arrays.append(
162+
[weakref.ref(c.keep[0]) for c in out[1]] if out is not None else []
163+
)
164+
return out
165+
166+
monkeypatch.setattr(columnar, "pack", pack)
167+
168+
169+
def _vector_type(db, kind="VERTEX"):
170+
db.command("sql", f"CREATE {kind} TYPE V")
171+
db.command("sql", "CREATE PROPERTY V.id LONG")
172+
db.command("sql", "CREATE PROPERTY V.embedding ARRAY_OF_FLOATS")
173+
174+
175+
@pytest.mark.parametrize("kind", ["VERTEX", "DOCUMENT"])
176+
def test_every_batch_is_packed_once_and_released_before_the_next(
177+
temp_db, monkeypatch, kind
178+
):
179+
_vector_type(temp_db, kind)
180+
spy = _PackSpy(monkeypatch)
181+
n, batch = 2_505, 1_000
182+
table = pa.table(
183+
{"id": pa.array(np.arange(n, dtype=np.int64)), "embedding": _fixed(_vectors(n))}
184+
)
185+
assert temp_db.insert_arrow("V", table, batch_size=batch) == n
186+
assert spy.rows == [
187+
1_000,
188+
1_000,
189+
505,
190+
] # one pack per batch, never more than batch_size rows
191+
assert spy.alive_at_next_call == [
192+
0,
193+
0,
194+
0,
195+
], spy.alive_at_next_call # batch k is gone when batch k + 1 is packed
196+
gc.collect()
197+
assert not any(
198+
ref() is not None for refs in spy.arrays for ref in refs
199+
) # and the last one when the call returns
200+
201+
202+
def test_a_dataframe_load_is_packed_per_batch_too(temp_db, monkeypatch):
203+
import pandas as pd
204+
205+
temp_db.command("sql", "CREATE VERTEX TYPE P")
206+
temp_db.command("sql", "CREATE PROPERTY P.id LONG")
207+
temp_db.command("sql", "CREATE PROPERTY P.name STRING")
208+
spy = _PackSpy(monkeypatch)
209+
frame = pd.DataFrame(
210+
{"id": np.arange(250, dtype=np.int64), "name": [f"p{i}" for i in range(250)]}
211+
)
212+
assert temp_db.insert_dataframe("P", frame, batch_size=100) == 250
213+
assert spy.rows == [100, 100, 50] and spy.alive_at_next_call == [0, 0, 0]
214+
215+
216+
def test_the_python_side_peak_of_a_whole_vector_load_stays_below_one_batch_of_python_floats(
217+
temp_db,
218+
):
219+
_vector_type(temp_db)
220+
n, batch = 40_000, 10_000
221+
table = pa.table(
222+
{"id": pa.array(np.arange(n, dtype=np.int64)), "embedding": _fixed(_vectors(n))}
223+
)
224+
one_batch_of_python_floats = _peak(
225+
lambda: table.slice(0, batch).column("embedding").to_pylist()
226+
)
227+
whole_load = _peak(lambda: temp_db.insert_arrow("V", table, batch_size=batch))
228+
assert one_batch_of_python_floats > 10 * 1024 * 1024
229+
assert whole_load * 2 < one_batch_of_python_floats, (
230+
whole_load,
231+
one_batch_of_python_floats,
232+
) # measured: about 7 times less
233+
assert temp_db.query("sql", "SELECT count(*) AS n FROM V").first().get("n") == n
234+
235+
236+
def _stored(db, key):
237+
return [
238+
float(x)
239+
for x in db.query("sql", "SELECT embedding FROM V WHERE id = ?", key)
240+
.first()
241+
.get("embedding")
242+
]
243+
244+
245+
@pytest.mark.parametrize("kind", ["VERTEX", "DOCUMENT"])
246+
def test_the_stored_vectors_are_the_ones_the_python_lists_stored(temp_db, kind):
247+
_vector_type(temp_db, kind)
248+
mat = _vectors(120)
249+
table = pa.table(
250+
{"id": pa.array(np.arange(120, dtype=np.int64)), "embedding": _fixed(mat)}
251+
)
252+
assert (
253+
temp_db.insert_arrow("V", table.slice(5, 100), batch_size=30) == 100
254+
) # a sliced table, four batches
255+
for key in (5, 6, 34, 35, 64, 104):
256+
assert (
257+
max(abs(a - b) for a, b in zip(_stored(temp_db, key), mat[key].tolist()))
258+
< 1e-6
259+
)
260+
assert temp_db.query("sql", "SELECT count(*) AS n FROM V").first().get("n") == 100
261+
262+
263+
def test_a_vector_column_with_a_null_or_a_list_property_takes_the_row_path_with_the_same_values(
264+
temp_db,
265+
):
266+
_vector_type(temp_db)
267+
temp_db.command("sql", "CREATE VERTEX TYPE L")
268+
temp_db.command("sql", "CREATE PROPERTY L.id LONG")
269+
temp_db.command("sql", "CREATE PROPERTY L.tags LIST")
270+
with_null = pa.table(
271+
{
272+
"id": pa.array([1, 2, 3], type=pa.int64()),
273+
"embedding": pa.array(
274+
[[0.5, 0.25], None, [1.0, 2.0]], type=pa.list_(pa.float32())
275+
),
276+
}
277+
)
278+
assert temp_db.insert_arrow("V", with_null) == 3
279+
assert [
280+
float(x)
281+
for x in temp_db.query("sql", "SELECT embedding FROM V WHERE id = 1")
282+
.first()
283+
.get("embedding")
284+
] == [0.5, 0.25]
285+
assert (
286+
temp_db.query("sql", "SELECT embedding FROM V WHERE id = 2")
287+
.first()
288+
.get("embedding")
289+
is None
290+
)
291+
tags = pa.table(
292+
{
293+
"id": pa.array([1, 2], type=pa.int64()),
294+
"tags": pa.array([[1, 2, 3], [4]], type=pa.list_(pa.int64())),
295+
}
296+
)
297+
assert temp_db.insert_arrow("L", tags) == 2
298+
assert list(
299+
temp_db.query("sql", "SELECT tags FROM L WHERE id = 1").first().get("tags")
300+
) == [1, 2, 3]

0 commit comments

Comments
 (0)