diff --git a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java index 9b5acb91..1703b965 100644 --- a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java +++ b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java @@ -561,7 +561,8 @@ public static CCharPointer insertCols(final IsolateThread thread, final long h, * Columnar insert by pointer: no copy of the columns into one buffer. spec points at 6 longs per column (data address and bytes, aux * address and bytes, null-bitmap address and bytes, see BulkOps.insertColumns) and nspec is their count. Kinds q d b f as * arcadedb_n_insert_cols, plus i (int32), e (float32) and s (Arrow UTF-8 string: int32 offsets and the bytes they index), and null bitmaps. - * The caller keeps every buffer alive and unchanged until the call returns. Returns {"ok":true,"n":rows}. + * The caller keeps every buffer alive and unchanged until the call returns. Returns {"ok":true,"n":rows}; with "parallel":true in the + * header the records go to the async executor and the reply also carries "fid" (see BulkOps.insertColumnsReply). */ @CEntryPoint(name = "arcadedb_n_insert_colp") public static CCharPointer insertColp(final IsolateThread thread, final long h, final CCharPointer type, final CCharPointer header, @@ -572,7 +573,7 @@ public static CCharPointer insertColp(final IsolateThread thread, final long h, final long[] table = new long[nspec]; for (int i = 0; i < nspec; i++) table[i] = ((Pointer) spec).readLong(WordFactory.signed(8L * i)); - return out("{\"ok\":true,\"n\":" + BulkOps.insertColumns(db(h), str(type), str(header), table) + "}"); + return out(BulkOps.insertColumnsReply(db(h), str(type), str(header), table)); } catch (final Throwable t) { return err(t); } diff --git a/native/src/main/java/org/humemai/nativelib/BulkOps.java b/native/src/main/java/org/humemai/nativelib/BulkOps.java index 37cedf90..cdf3cd98 100644 --- a/native/src/main/java/org/humemai/nativelib/BulkOps.java +++ b/native/src/main/java/org/humemai/nativelib/BulkOps.java @@ -10,6 +10,9 @@ import com.arcadedb.database.Database; import com.arcadedb.database.MutableDocument; +import com.arcadedb.database.async.DatabaseAsyncExecutor; +import com.arcadedb.database.async.ErrorCallback; +import com.arcadedb.schema.EdgeType; import com.arcadedb.schema.VertexType; import org.graalvm.nativeimage.c.type.CTypeConversion; @@ -55,13 +58,25 @@ static long insertColumns(final Database db, final String type, final String hea spec[SLOTS * c + 1] = bytes; pos += bytes; } - return insertColumns(db, type, h, spec); + return insertColumns(db, type, h, spec, null); } - /** The pointer form (arcadedb_n_insert_colp): the table holds SLOTS longs per column, see {@link #insertColumns(Database, String, Map, long[])}. */ + /** + * The pointer form (arcadedb_n_insert_colp): the table holds SLOTS longs per column, see {@link #insertColumns(Database, String, Map, long[], ErrorCallback)}. + * Returns the reply text {"ok":true,"n":rows}. A header with "parallel":true must also carry "fid", a failure token that the caller made + * with the "token" operation of arcadedb_n_async: the records go to the database's async executor, the call does not wait for the writers + * and neither opens, commits nor rolls back a transaction ("commit" does not apply), exactly as insert_parallel does for rows. The + * caller waits for the writers and reads (and so removes) the token with the "failures" operation, also when this call throws or never + * returns to it: the token is the caller's, so this call never drops it. + */ @SuppressWarnings("unchecked") - static long insertColumns(final Database db, final String type, final String header, final long[] spec) { - return insertColumns(db, type, (Map) Json.parse(header), spec); + static String insertColumnsReply(final Database db, final String type, final String header, final long[] spec) { + final Map h = (Map) Json.parse(header); + if (!Boolean.TRUE.equals(h.get("parallel"))) + return "{\"ok\":true,\"n\":" + insertColumns(db, type, h, spec, null) + "}"; + if (!(h.get("fid") instanceof Number fid)) + throw new IllegalArgumentException("a parallel columnar insert needs a failure token (header \"fid\", from the \"token\" operation)"); + return "{\"ok\":true,\"n\":" + insertColumns(db, type, h, spec, NativeExtras.FailureToken.of(fid.longValue()).onError) + "}"; } /** @@ -72,9 +87,13 @@ static long insertColumns(final Database db, final String type, final String hea * means no nulls. The memory belongs to the caller for the duration of the call; nothing is kept. Every length is checked here before * the first document is built, so a wrong table is an exception and nothing is written. Same transaction contract as * DocumentBatcher.insertColumns: batch commits and the rollback on failure apply only to a transaction opened here. + * onError null: the synchronous mode above. Not null (header "parallel":true): every record goes to db.async().createRecord with that + * callback and no transaction is touched, as DocumentBatcher.insertColumnsParallel does; an edge type is refused there, as it is by + * insert_parallel. The columns are validated before the first record is handed over in both modes. */ @SuppressWarnings("unchecked") - private static long insertColumns(final Database db, final String type, final Map h, final long[] spec) { + private static long insertColumns(final Database db, final String type, final Map h, final long[] spec, + final ErrorCallback onError) { final long nRows = ((Number) h.get("n")).longValue(); if (nRows < 0 || nRows > Integer.MAX_VALUE) throw new IllegalArgumentException("row count out of range: " + nRows); @@ -137,45 +156,33 @@ private static long insertColumns(final Database db, final String type, final Ma } } final byte[] scratch = new byte[longest]; - final boolean vertex = db.getSchema().getType(type) instanceof VertexType; + final Object schemaType = db.getSchema().getType(type); + final boolean vertex = schemaType instanceof VertexType; + if (onError != null) { + if (schemaType instanceof EdgeType) + throw new IllegalArgumentException("Cannot bulk-insert into '" + type + + "' because it is an edge type; load the vertices, then create the edges (graph_batch().new_edges() or CREATE EDGE)"); + final DatabaseAsyncExecutor async = db.async(); + int handed = 0; + try { + for (; handed < n; handed++) { + final MutableDocument doc = vertex ? db.newVertex(type) : db.newDocument(type); + setRow(doc, handed, names, kinds, dims, data, nulls, blobs, scratch); + async.createRecord(doc, null, onError); + } + } catch (final Throwable t) { + // the records handed over before the failure are the writers' now, whatever this call does + throw new IllegalStateException("after handing " + handed + " of " + n + " rows to the async writers (those may have been stored): " + t, t); + } + return n; + } final boolean wasActive = db.isTransactionActive(); if (!wasActive) db.begin(); try { for (int i = 0; i < n; i++) { final MutableDocument doc = vertex ? db.newVertex(type) : db.newDocument(type); - for (int c = 0; c < k; c++) { - final Pointer base = WordFactory.pointer(data[c]); - if (nulls[c] != 0 && ((WordFactory.pointer(nulls[c]).readByte(WordFactory.signed(i >> 3)) >> (i & 7)) & 1) == 0) { - doc.set(names[c], null); - continue; - } - switch (kinds[c]) { - case 'q' -> doc.set(names[c], base.readLong(WordFactory.signed(8L * i))); - case 'd' -> doc.set(names[c], base.readDouble(WordFactory.signed(8L * i))); - case 'b' -> doc.set(names[c], base.readByte(WordFactory.signed(i)) != 0); - case 'i' -> doc.set(names[c], base.readInt(WordFactory.signed(4L * i))); - case 'e' -> doc.set(names[c], base.readFloat(WordFactory.signed(4L * i))); - case 's' -> { - final int from = base.readInt(WordFactory.signed(4L * i)); - final int len = base.readInt(WordFactory.signed(4L * (i + 1L))) - from; - if (len == 0) - doc.set(names[c], ""); - else { - blobs[c].get(from, scratch, 0, len); - doc.set(names[c], new String(scratch, 0, len, StandardCharsets.UTF_8)); - } - } - default -> { - final int dim = dims[c]; - final float[] v = new float[dim]; - final long at = 4L * i * dim; - for (int j = 0; j < dim; j++) - v[j] = base.readFloat(WordFactory.signed(at + 4L * j)); - doc.set(names[c], v); - } - } - } + setRow(doc, i, names, kinds, dims, data, nulls, blobs, scratch); doc.save(); if (!wasActive && commitEvery > 0 && (i + 1) % commitEvery == 0) { db.commit(); @@ -197,6 +204,44 @@ private static long insertColumns(final Database db, final String type, final Ma return n; } + /** Sets the properties of row i on doc from the columns: the one place that reads a value out of a buffer. */ + private static void setRow(final MutableDocument doc, final int i, final String[] names, final char[] kinds, final int[] dims, + final long[] data, final long[] nulls, final ByteBuffer[] blobs, final byte[] scratch) { + final int k = names.length; + for (int c = 0; c < k; c++) { + final Pointer base = WordFactory.pointer(data[c]); + if (nulls[c] != 0 && ((WordFactory.pointer(nulls[c]).readByte(WordFactory.signed(i >> 3)) >> (i & 7)) & 1) == 0) { + doc.set(names[c], null); + continue; + } + switch (kinds[c]) { + case 'q' -> doc.set(names[c], base.readLong(WordFactory.signed(8L * i))); + case 'd' -> doc.set(names[c], base.readDouble(WordFactory.signed(8L * i))); + case 'b' -> doc.set(names[c], base.readByte(WordFactory.signed(i)) != 0); + case 'i' -> doc.set(names[c], base.readInt(WordFactory.signed(4L * i))); + case 'e' -> doc.set(names[c], base.readFloat(WordFactory.signed(4L * i))); + case 's' -> { + final int from = base.readInt(WordFactory.signed(4L * i)); + final int len = base.readInt(WordFactory.signed(4L * (i + 1L))) - from; + if (len == 0) + doc.set(names[c], ""); + else { + blobs[c].get(from, scratch, 0, len); + doc.set(names[c], new String(scratch, 0, len, StandardCharsets.UTF_8)); + } + } + default -> { + final int dim = dims[c]; + final float[] v = new float[dim]; + final long at = 4L * i * dim; + for (int j = 0; j < dim; j++) + v[j] = base.readFloat(WordFactory.signed(at + 4L * j)); + doc.set(names[c], v); + } + } + } + } + /** * Documents (or vertices of a vertex type) from rows [{property: value},..] in the UTF-8 JSON text b[0, n): one record per row through * newDocument/newVertex, set per property and save, as DocumentBatcher.insertManyJson does on the JPype path; no transaction is opened diff --git a/native/src/main/java/org/humemai/nativelib/NativeExtras.java b/native/src/main/java/org/humemai/nativelib/NativeExtras.java index a62574ad..404181b5 100644 --- a/native/src/main/java/org/humemai/nativelib/NativeExtras.java +++ b/native/src/main/java/org/humemai/nativelib/NativeExtras.java @@ -527,6 +527,51 @@ else if (opName.equals("create_vertices")) private static final Map ASYNC_FAILS = new ConcurrentHashMap<>(); private static final Map> ASYNC_FIRST = new ConcurrentHashMap<>(); + /** + * A failure token for one parallel load: the id the caller reads (and so removes) with the "failures" operation, and the callback that + * counts. The caller creates it before the load ("token" operation) and passes the id in the load's header, so that it holds the id + * whatever happens to the call, an interrupt included, and always has the means to wait for the writers and remove the entries. + */ + static final class FailureToken { + final long fid; + final AtomicLong count; + final AtomicReference first; + final ErrorCallback onError; + + private FailureToken(final long fid, final AtomicLong count, final AtomicReference first) { + this.fid = fid; + this.count = count; + this.first = first; + this.onError = t -> { + count.incrementAndGet(); + first.compareAndSet(null, t); + }; + } + + /** A new token, registered for the "failures" operation. */ + static FailureToken create() { + final FailureToken token = new FailureToken(ArcadeNative.IDS.getAndIncrement(), new AtomicLong(), new AtomicReference<>()); + ASYNC_FAILS.put(token.fid, token.count); + ASYNC_FIRST.put(token.fid, token.first); + return token; + } + + /** The token that "token" created, for the load that its id was passed to. */ + static FailureToken of(final long fid) { + final AtomicLong count = ASYNC_FAILS.get(fid); + final AtomicReference first = ASYNC_FIRST.get(fid); + if (count == null || first == null) + throw new IllegalArgumentException("unknown failure token " + fid + " (already read, or never created)"); + return new FailureToken(fid, count, first); + } + + /** The load failed before it returned: the caller gets the error, not a token to read failures from. */ + void drop() { + ASYNC_FAILS.remove(fid); + ASYNC_FIRST.remove(fid); + } + } + private static String replyRows(final ResultSet rs) { final StringBuilder sb = new StringBuilder("{\"ok\":true,\"rs\":0,"); NativeCodec.appendBatch(sb, rs, Integer.MAX_VALUE); @@ -824,15 +869,9 @@ else if (a.get("slot") instanceof Number slot) ex.appendSamples((String) a.get("type"), stamps, columns); } case "insert_parallel" -> { - final long fid = ArcadeNative.IDS.getAndIncrement(); - final AtomicLong count = new AtomicLong(); - final AtomicReference first = new AtomicReference<>(); - ASYNC_FAILS.put(fid, count); - ASYNC_FIRST.put(fid, first); - final ErrorCallback onError = t -> { - count.incrementAndGet(); - first.compareAndSet(null, t); - }; + final FailureToken token = FailureToken.create(); + final long fid = token.fid; + final ErrorCallback onError = token.onError; final String type = (String) a.get("type"); try { final List rows = (List) a.get("rows"); @@ -846,12 +885,14 @@ else if (a.get("slot") instanceof Number slot) ex.createRecord(doc, null, onError); } } catch (final Throwable t) { - ASYNC_FAILS.remove(fid); // the caller gets the error, not a token to read failures from - ASYNC_FIRST.remove(fid); + token.drop(); // the caller gets the error, not a token to read failures from throw t; } return ArcadeNative.out("{\"ok\":true,\"fid\":" + fid + "}"); } + case "token" -> { + return ArcadeNative.out("{\"ok\":true,\"fid\":" + FailureToken.create().fid + "}"); + } case "failures" -> { final long fid = ((Number) a.get("fid")).longValue(); final AtomicLong count = ASYNC_FAILS.remove(fid); diff --git a/src/arcadedb/_native/bulk.py b/src/arcadedb/_native/bulk.py index d244c48c..9851848f 100644 --- a/src/arcadedb/_native/bulk.py +++ b/src/arcadedb/_native/bulk.py @@ -2,11 +2,13 @@ ints, floats and bools, lists of one value type (None is a null), strings (UTF-8 bytes plus offsets, as Arrow lays them out), float32 vector columns and Arrow arrays are laid out by :mod:`.columnar`; anything else (nested values, mixed-type columns, date-times, decimals, integers past 64 bits) returns None and the caller uses the per-row JSON path. The documents are built image-side as -``DocumentBatcher.insertColumns`` builds them on JPype. ``ARCADEDB_NATIVE_BULK=0`` turns the columnar path off. +``DocumentBatcher.insertColumns`` builds them on JPype; with ``parallel`` they are handed to the async writers instead of a transaction +(``DocumentBatcher.insertColumnsParallel``). ``ARCADEDB_NATIVE_BULK=0`` turns the columnar path off. """ import os +from ..exceptions import ArcadeDBError from . import columnar, lib _ENABLED = os.environ.get("ARCADEDB_NATIVE_BULK", "1") != "0" @@ -21,20 +23,19 @@ def _usable(): ) -def _send(db, type_name, packed, rows, commit_every, prefix): +def _send(db, type_name, packed, rows, commit_every, prefix, fid=None): from .database import _call spec = columnar.spec_of(packed) - reply = _call( + return _call( prefix, "arcadedb_n_insert_colp", db._h, lib.b(type_name), - lib.b(columnar.header_of(rows, commit_every, packed)), + lib.b(columnar.header_of(rows, commit_every, packed, fid)), spec, len(spec), ) # `packed` holds the arrays the addresses point into until the call returns - return int(reply["n"]) def insert_columns_fast( @@ -51,7 +52,32 @@ def insert_columns_fast( if packed is None: return None rows, cols = packed - return _send(db, type_name, cols, rows, commit_every, prefix) + return int(_send(db, type_name, cols, rows, commit_every, prefix)["n"]) + + +def insert_columns_parallel(db, type_name, columns, declared=None): + """Hand ``columns`` (name -> column) to the async writers in one crossing, wait for them and return the row count, or None when a + column cannot be sent this way (nothing was handed over then). A record that a writer rejects raises + (``Database._finish_parallel``), as for ``insert_many(parallel=True)``, and so does an interrupt, after the writers have finished: + the failure token is made before the crossing, so the call always has the id to wait for and to read. + """ + if not _usable(): + return None + packed = columnar.pack(columns, columnar.Declared(db, type_name, declared)) + if packed is None: + return None + rows, cols = packed + from . import extras + + fid = extras.failure_token(db) + try: + _send(db, type_name, cols, rows, 0, "Failed to bulk-insert", fid=fid) + except ArcadeDBError as e: + raise db._parallel_load_failed(type_name, e, [fid]) from e + except BaseException: + db._abandon_parallel([fid]) # a ^C: the writers still run, so wait for them + raise + return db._finish_parallel(type_name, [fid], rows) def insert_rows_columnar(db, type_name, chunk): diff --git a/src/arcadedb/_native/columnar.py b/src/arcadedb/_native/columnar.py index f288043c..27d261a4 100644 --- a/src/arcadedb/_native/columnar.py +++ b/src/arcadedb/_native/columnar.py @@ -343,17 +343,22 @@ def spec_of(packed): return (ctypes.c_int64 * len(flat))(*flat) -def header_of(rows, commit_every, packed): - """The JSON header text: row count, commit interval and ``[name, kind, dim]`` per column.""" +def header_of(rows, commit_every, packed, fid=None): + """The JSON header text: row count, commit interval and ``[name, kind, dim]`` per column; with a failure token ``fid`` it also says + ``"parallel": true``, which hands the records to the async executor (the image then ignores the commit interval and counts the + failed records on that token). + """ import json - return json.dumps( - { - "n": rows, - "commit": min(int(commit_every), 2**31 - 1), - "cols": [[c.name, c.kind, c.dim] for c in packed], - } - ) + header = { + "n": rows, + "commit": min(int(commit_every), 2**31 - 1), + "cols": [[c.name, c.kind, c.dim] for c in packed], + } + if fid is not None: + header["parallel"] = True + header["fid"] = fid + return json.dumps(header) def rows_to_columns(rows): diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index f30dc526..6fab6c58 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -995,20 +995,50 @@ def _command_many(self, language: str, command: str, rows) -> int: raise _error("Command failed", reply) return reply["n"] - def _parallel_load_failed(self, type_name, error) -> ArcadeDBError: - """The error for a parallel load that failed while handing rows to the writers: waits for the queue to drain first, so - the caller sees a final state, and says that the rows handed over before the failure may have been stored. + def _drain_parallel(self, tokens): + """Waits for the async writers, then reads, and so removes, the failure token of every parallel load handed in. Returns + ``(problem, failed, first)``: the error of the wait or of a read (or None), the failed records the writers reported on all the + tokens and the first failure's text. The tokens are read whatever the wait did: an unread token keeps its entries for good. """ + from . import extras + + problem = None try: self._async().wait_completion() - except Exception: # nosec B110 - the original error is the one to report - pass + except Exception as e: + problem = e + failed, first = 0, None + for token in tokens: + try: + result = extras.failures(self, token) + except Exception as e: + problem = problem or e + continue + if result["count"] and first is None: + first = result["first"] + failed += result["count"] + return problem, failed, first + + def _abandon_parallel(self, tokens): + """A parallel load that is being interrupted: waits for the writers (they run on) and removes the tokens; the error that is on its + way out is the one to report. + """ + self._drain_parallel(tokens) + + def _parallel_load_failed(self, type_name, error, tokens=()) -> ArcadeDBError: + """The error for a parallel load that failed while handing rows to the writers: waits for the queue to drain first, so + the caller sees a final state, and says that the rows handed over before the failure may have been stored. + """ + _, failed, _ = self._drain_parallel(tokens) text = str(error) prefix = "Failed to bulk-insert: " if text.startswith(prefix): text = text[len(prefix) :] + reported = ( + f"; the writers had reported {failed} failed record(s)" if failed else "" + ) return ArcadeDBError( - f"Failed to bulk-insert into '{type_name}': {text}; the rows handed to the " + f"Failed to bulk-insert into '{type_name}': {text}{reported}; the rows handed to the " f"parallel writers before the failure may have been stored (the call waited " f"for them before raising)" ) @@ -1030,26 +1060,28 @@ def _insert_many_parallel(self, type_name, rows) -> int: tokens.append(extras.insert_parallel(self, type_name, chunk)["fid"]) except ArcadeDBError as e: # as the JPype path: the rows handed to the writers before the failure are written whatever this call does - raise self._parallel_load_failed(type_name, e) from e + raise self._parallel_load_failed(type_name, e, tokens) from e n += len(chunk) except ArcadeDBError: raise except BaseException: - try: - self._async().wait_completion() - except Exception: # nosec B110 - the original error is the one to report - pass + self._abandon_parallel(tokens) raise + return self._finish_parallel(type_name, tokens, n) + + def _finish_parallel(self, type_name, tokens, n) -> int: + """The end of a parallel load: wait for the async writers, then read the failure count of every token and raise when a record failed.""" try: - self._async().wait_completion() - except Exception as e: - raise ArcadeDBError(f"Failed to bulk-insert into '{type_name}': {e}") from e - n_failed, first = 0, None - for token in tokens: - result = extras.failures(self, token) - if result["count"] and first is None: - first = result["first"] - n_failed += result["count"] + problem, n_failed, first = self._drain_parallel(tokens) + except BaseException: + self._abandon_parallel( + tokens + ) # a ^C in the wait: finish it, and remove the tokens + raise + if problem is not None: + raise ArcadeDBError( + f"Failed to bulk-insert into '{type_name}': {problem}" + ) from problem if n_failed: raise ArcadeDBError( f"Failed to bulk-insert into '{type_name}': the parallel writers reported {n_failed} failed record(s) " @@ -1314,7 +1346,11 @@ def export_to_csv( def insert_columns( self, type_name, columns, commit_every: int = 10_000, parallel: bool = False ) -> int: - """Bulk insert from whole columns (dict of name -> numpy array or list), as rows through :meth:`insert_many`.""" + """Bulk insert from whole columns (dict of name -> numpy array or list): columns the buffers can carry cross in one call + (``arcadedb_n_insert_colp``), any other column set goes as rows through :meth:`insert_many`. With ``parallel=True`` the + same one call hands the records to the async writers (no transaction, ``commit_every`` does not apply, the call waits for the + writers and raises when a record failed), as ``insert_many(parallel=True)`` does for rows. + """ self._check_not_closed() items = list(columns.items()) if hasattr(columns, "items") else None if not items: @@ -1330,7 +1366,13 @@ def insert_columns( known = {n: len(v) for n, v in items if hasattr(v, "__len__")} if len(set(known.values())) > 1: raise ValueError(f"columns differ in length: {known}") - if not parallel: + if parallel: + from .bulk import insert_columns_parallel + + done = insert_columns_parallel(self, type_name, columns) + if done is not None: + return done + else: from .bulk import insert_columns_fast done = insert_columns_fast(self, type_name, columns, commit_every) diff --git a/src/arcadedb/_native/extras.py b/src/arcadedb/_native/extras.py index 5f8a4702..439fb9ef 100644 --- a/src/arcadedb/_native/extras.py +++ b/src/arcadedb/_native/extras.py @@ -915,6 +915,19 @@ def insert_parallel(database, type_name: str, rows) -> dict: ) +def failure_token(database) -> int: + """A new failure token for one parallel load: the caller makes it before the load, hands its id to the load and reads it with + :func:`failures` after the final wait, so it holds the id whatever happens to the load (an interrupt included). + """ + return _call( + "Failed to make a failure token", + "arcadedb_n_async", + database._h, + b"token", + lib.b(_enc({})), + )["fid"] + + def failures(database, token: int) -> dict: return _call( "Failed to read the parallel writers", diff --git a/src/arcadedb/dataframe.py b/src/arcadedb/dataframe.py index 9d3ee0b7..46362e53 100644 --- a/src/arcadedb/dataframe.py +++ b/src/arcadedb/dataframe.py @@ -421,8 +421,12 @@ def insert_table( native = _is_native(db) written = 0 # a document load inside the caller's own transaction joins it (insert_columns commits - # nothing then); a vertex load refuses to run inside one - in_caller_transaction = (not is_vertex) and bool(db.is_transaction_active()) + # nothing then); a vertex load refuses to run inside one. parallel=True never joins it: the + # async writers commit the records themselves, so a rollback of the caller's transaction does + # not remove them + in_caller_transaction = ( + (not is_vertex) and (not parallel) and bool(db.is_transaction_active()) + ) graph_batch = db.graph_batch() if is_vertex else None try: for batch in table.to_batches(max_chunksize=batch_size): @@ -434,7 +438,7 @@ def insert_table( batch.column(i), pa, native, - arrow=native and not is_vertex and not parallel, + arrow=native and not is_vertex, ) for i, name in enumerate(names) } @@ -459,6 +463,13 @@ def insert_table( f"{exc}; the {written} row(s) of the batches before this one " f"were written in your transaction, not committed" ) from exc + if written and parallel: + # the async writers commit on their own, inside the caller's transaction or not + raise ArcadeDBError( + f"{exc}; the {written} row(s) of the batches before this one " + f"were handed to the async writers, which commit on their own: " + f"a rollback of your transaction does not remove them" + ) from exc if written: raise ArcadeDBError( f"{exc}; the {written} row(s) of the batches before this one " @@ -497,9 +508,20 @@ def _insert_documents_native( """One batch of documents on the native backend. When every column can cross as memory (numbers, bools, strings, nulls, float vectors for an ``ARRAY_OF_FLOATS`` property) the batch goes through the columnar entry, no Python values and no JSON text; otherwise the columns become one dict per row and go to ``insert_many`` (one crossing per 10,000 rows). Either way - the batch is one transaction, or joins the caller's. + the batch is one transaction, or joins the caller's. With ``parallel`` the same columnar entry hands the records to the async + writers (no transaction), and the call waits for them and raises when a record failed, as ``insert_many(parallel=True)`` does. """ - if not parallel: + if parallel: + from ._native.bulk import insert_columns_parallel + + if ( + insert_columns_parallel( + db, type_name, columns, {} if declared is None else declared + ) + is not None + ): + return + else: from ._native.bulk import insert_columns_fast if ( diff --git a/tests/conftest.py b/tests/conftest.py index ec3147c7..707833d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -47,6 +47,10 @@ collect_ignore.append("test_native_vector_hit_lazy.py") # the columnar insert's stored records are compared between the JSON and the columnar entry points of the native library collect_ignore.append("test_native_columnar_insert.py") + # insert_columns(parallel=True) takes the columnar entry and stores what the row path stores (arcadedb_n_insert_colp with "parallel") + collect_ignore.append("test_native_columnar_parallel.py") + # the review of the native parallel loads: no silent loss, an interrupt waits for the writers and reads the failure token + collect_ignore.append("test_native_columnar_parallel_review_fixes.py") # the rows of an insert_many chunk with date-times, Decimals, bytes and numpy scalars are told apart by the entry points they cross collect_ignore.append("test_native_mixed_chunk.py") # the native build's ingest defaults (insert_many and graph_batch commit size), observed on the native library diff --git a/tests/test_native_columnar_pack.py b/tests/test_native_columnar_pack.py index cda0a1f7..af26f211 100644 --- a/tests/test_native_columnar_pack.py +++ b/tests/test_native_columnar_pack.py @@ -233,6 +233,11 @@ def test_the_address_table_and_header_describe_the_columns(): "cols": [["id", "q", 0], ["s", "s", 0]], } assert json.loads(columnar.header_of(rows, 2**40, packed))["commit"] == 2**31 - 1 + # "parallel" and the failure token are in the header only when a token is given (the image hands the records to the async writers then) + parallel = json.loads(columnar.header_of(rows, 10_000, packed, fid=41)) + assert parallel["parallel"] is True and parallel["fid"] == 41 + plain = json.loads(columnar.header_of(rows, 10_000, packed)) + assert "parallel" not in plain and "fid" not in plain def test_rows_become_columns_only_when_every_row_has_the_same_keys(): diff --git a/tests/test_native_columnar_parallel.py b/tests/test_native_columnar_parallel.py new file mode 100644 index 00000000..90d6f79a --- /dev/null +++ b/tests/test_native_columnar_parallel.py @@ -0,0 +1,242 @@ +"""Native ``insert_columns(parallel=True)`` takes the columnar entry and stores what the row path of ``insert_many(parallel=True)`` stores. + +Before, a parallel call skipped the columnar fast path: the columns became one dict per row, then one JSON text per chunk of 10,000 +rows for ``arcadedb_n_async insert_parallel``. Now the same ``arcadedb_n_insert_colp`` crossing that the synchronous call uses carries +``"parallel": true`` and the image hands the records to the async writers, as ``DocumentBatcher.insertColumnsParallel`` does on the +JVM wheel. Needs the native library, so conftest.py collects this file on the native backend only. Each test loads the same data +through the row path (the private reference: ``insert_many(parallel=True)`` with dict rows) and through the columns, compares the +stored records value by value and type by type, and spies on the entry points so a silent fall back cannot pass. +""" + +import numpy as np +import pytest + +from arcadedb._native import bulk, extras, lib # noqa: E402 +from arcadedb.exceptions import ArcadeDBError # noqa: E402 + +COLP, ROWS = "arcadedb_n_insert_colp", "arcadedb_n_insert_rows" + + +class Spy: + """Counts the crossings by name and the calls that hand a chunk of dict rows to the parallel writers.""" + + def __init__(self, monkeypatch): + self.calls, self.row_chunks = [], 0 + real_call, real_rows = lib.call, extras.insert_parallel + + def call(name, *args): + self.calls.append(name) + return real_call(name, *args) + + def insert_parallel(database, type_name, rows): + self.row_chunks += 1 + return real_rows(database, type_name, rows) + + monkeypatch.setattr(lib, "call", call) + monkeypatch.setattr(extras, "insert_parallel", insert_parallel) + + def count(self, name): + return self.calls.count(name) + + +@pytest.fixture +def spy(monkeypatch): + return Spy(monkeypatch) + + +def _count(db, name): + sql = f"SELECT count(*) AS n FROM {name}" # nosec B608 - test-owned type name + return db.query("sql", sql).to_list()[0]["n"] + + +def _types(db, *names, kind="DOCUMENT", buckets=3, unique=False): + for name in names: + db.command("sql", f"CREATE {kind} TYPE {name} BUCKETS {buckets}") + db.command("sql", f"CREATE PROPERTY {name}.id INTEGER") + db.command("sql", f"CREATE PROPERTY {name}.val DOUBLE") + db.command("sql", f"CREATE PROPERTY {name}.flag BOOLEAN") + db.command("sql", f"CREATE PROPERTY {name}.txt STRING") + db.command("sql", f"CREATE PROPERTY {name}.vec ARRAY_OF_FLOATS") + if unique: + db.command("sql", f"CREATE INDEX ON {name} (id) UNIQUE") + + +def _stored(db, name): + """Every record of a type as (value, python type) pairs per property, ordered by id (the writers store in no fixed order).""" + sql = f"SELECT FROM {name} ORDER BY id" # nosec B608 - test-owned type name + return [ + {k: (v, type(v).__name__) for k, v in row.items() if not k.startswith("@")} + for row in db.query("sql", sql).to_list() + ] + + +def _columns(n, dim=4, seed=11): + rng = np.random.default_rng(seed) + ids = np.arange(n, dtype=np.int64) + return { + "id": ids, + "val": (ids * 37 % 10007) / 7.0, + "flag": ids % 3 == 0, + "txt": ["row-%d-%s" % (k, "é" * (k % 5)) for k in range(n)], + "vec": rng.standard_normal((n, dim), dtype=np.float32), + } + + +def _rows(columns): + names = list(columns) + lists = [c.tolist() if hasattr(c, "tolist") else c for c in columns.values()] + return [dict(zip(names, vals)) for vals in zip(*lists)] + + +def test_a_parallel_call_uses_the_columnar_entry_and_stores_the_row_paths_records( + temp_db, spy +): + _types(temp_db, "Ref", "Col") + columns = _columns(3_100) + assert temp_db.insert_many("Ref", _rows(columns), parallel=True) == 3_100 + assert spy.row_chunks == 1 and spy.count(COLP) == 0 # the reference is the row path + + spy.calls.clear() + spy.row_chunks = 0 + assert temp_db.insert_columns("Col", columns, parallel=True) == 3_100 + assert spy.count(COLP) == 1, spy.calls # one crossing for the whole call + assert spy.row_chunks == 0 and spy.count(ROWS) == 0 # no dict rows, no JSON text + + expected = _stored(temp_db, "Ref") + assert len(expected) == 3_100 and _count(temp_db, "Col") == 3_100 + assert expected[7]["id"][1] == "int" and expected[7]["vec"][1] == "list" + assert _stored(temp_db, "Col") == expected + + +def test_nulls_and_awkward_strings_are_stored_as_the_row_path_stores_them(temp_db, spy): + _types(temp_db, "Ref", "Col") + columns = { + "id": list(range(8)), + "txt": ["", "a", None, "snow☃man", "x" * 300, None, "\n\t", "z"], + "val": [1.5, None, 2.5, 1e300, -1e-300, 0.0, 3.0, None], + } + temp_db.insert_many("Ref", _rows(columns), parallel=True) + spy.calls.clear() + temp_db.insert_columns("Col", columns, parallel=True) + assert spy.count(COLP) == 1 + got, want = _stored(temp_db, "Col"), _stored(temp_db, "Ref") + assert got == want and got[2]["txt"][0] is None + + +def test_a_vertex_type_gets_vertices_in_parallel(temp_db, spy): + _types(temp_db, "Ref", "Col", kind="VERTEX") + columns = _columns(1_500) + temp_db.insert_many("Ref", _rows(columns), parallel=True) + spy.calls.clear() + temp_db.insert_columns("Col", columns, parallel=True) + assert spy.count(COLP) == 1 + assert _stored(temp_db, "Col") == _stored(temp_db, "Ref") + kinds = temp_db.query( + "sql", "SELECT FROM schema:types WHERE name = 'Col'" + ).to_list() + assert kinds[0]["type"] == "vertex" + + +def test_a_record_the_writers_reject_is_reported_after_the_wait(temp_db, spy): + _types(temp_db, "Col", unique=True) + ids = np.arange(400, dtype=np.int64) + ids[300] = ids[7] # a duplicate key + with pytest.raises(ArcadeDBError, match="failed record"): + temp_db.insert_columns("Col", {"id": ids, "txt": ["t"] * 400}, parallel=True) + assert spy.count(COLP) == 1 and spy.row_chunks == 0 + stored = _count(temp_db, "Col") + assert 0 < stored <= 400 # the call had waited for the writers: nothing lands later + + +def test_a_value_the_declared_type_refuses_raises_and_the_call_is_final(temp_db, spy): + _types(temp_db, "Col") + values = [str(i) for i in range(1_000)] + values[700] = "not a number" + with pytest.raises( + ArcadeDBError, match="may have been stored|failed record" + ) as err: + temp_db.insert_columns("Col", {"id": values}, parallel=True) + assert "rolled back" not in str(err.value) + assert spy.count(COLP) == 1 + stored = _count(temp_db, "Col") + temp_db.async_executor().wait_completion() + assert _count(temp_db, "Col") == stored <= 1_000 + + +def test_an_edge_type_is_refused_as_before(temp_db): + temp_db.command("sql", "CREATE EDGE TYPE Link") + with pytest.raises(ArcadeDBError, match="edge type"): + temp_db.insert_columns("Link", {"id": [1, 2, 3]}, parallel=True) + + +def test_a_missing_type_is_the_schemas_error(temp_db): + with pytest.raises(ArcadeDBError, match="NoSuchType"): + temp_db.insert_columns("NoSuchType", {"id": [1, 2, 3]}, parallel=True) + + +def test_a_column_the_buffers_cannot_carry_still_goes_through_the_row_path( + temp_db, spy +): + _types(temp_db, "Col") + columns = {"id": [1, 2, 3, 4], "txt": ["a", 2, "c", 4.5]} # a mixed-type column + assert temp_db.insert_columns("Col", columns, parallel=True) == 4 + assert spy.count(COLP) == 0 and spy.row_chunks == 1 + assert _count(temp_db, "Col") == 4 + + +def test_the_switch_that_turns_the_columnar_path_off_applies_to_parallel( + temp_db, spy, monkeypatch +): + _types(temp_db, "Col") + monkeypatch.setattr(bulk, "_ENABLED", False) + assert temp_db.insert_columns("Col", _columns(50), parallel=True) == 50 + assert spy.count(COLP) == 0 and spy.row_chunks == 1 + assert _count(temp_db, "Col") == 50 + + +def test_no_rows_is_zero_and_writes_nothing(temp_db, spy): + _types(temp_db, "Col") + assert temp_db.insert_columns("Col", {"id": []}, parallel=True) == 0 + assert spy.count(COLP) == 0 and _count(temp_db, "Col") == 0 + + +def test_the_synchronous_call_is_unchanged(temp_db, spy): + _types(temp_db, "Col", unique=True) + assert temp_db.insert_columns("Col", _columns(500), commit_every=100) == 500 + assert spy.count(COLP) == 1 and spy.row_chunks == 0 + assert _count(temp_db, "Col") == 500 + + +def test_insert_arrow_parallel_takes_the_columnar_entry(temp_db, spy, monkeypatch): + import pyarrow as pa + + _types(temp_db, "Ref", "ByTable") + columns = _columns(2_000) + table = pa.table( + { + "id": columns["id"], + "val": pa.array( + [None if i % 7 == 0 else v for i, v in enumerate(columns["val"])] + ), + "txt": pa.array(columns["txt"]), + } + ) + with ( + monkeypatch.context() as m + ): # the private reference: the row path of the old code + m.setattr(bulk, "_ENABLED", False) + assert ( + temp_db.insert_arrow("Ref", table, batch_size=700, parallel=True) == 2_000 + ) + assert spy.count(COLP) == 0 and spy.row_chunks >= 1 + + spy.calls.clear() + spy.row_chunks = 0 + assert ( + temp_db.insert_arrow("ByTable", table, batch_size=700, parallel=True) == 2_000 + ) + assert spy.count(COLP) == 3 and spy.row_chunks == 0 # one crossing per batch + + expected = _stored(temp_db, "Ref") + assert len(expected) == 2_000 and expected[0]["val"][0] is None + assert _stored(temp_db, "ByTable") == expected diff --git a/tests/test_native_columnar_parallel_review_fixes.py b/tests/test_native_columnar_parallel_review_fixes.py new file mode 100644 index 00000000..70f17bc3 --- /dev/null +++ b/tests/test_native_columnar_parallel_review_fixes.py @@ -0,0 +1,203 @@ +"""Review of the native parallel loads (#328): nothing is lost silently, an interrupt leaves no writer running and no token behind. + +``insert_columns(parallel=True)`` hands the records to the async writers in one crossing and waits for them before it returns. These +tests pin what a caller relies on: a record that a writer rejects is counted (``stored + reported failed >= rows``) and no key is stored twice +(a writer that meets a duplicate abandons its whole open batch, and reports it); a load at the writers' batch boundary stores every row and nothing lands after the call returns; +a ^C during the crossing, or between the chunks of the row path, waits for the writers and reads (so removes) the failure token; a wait +that fails still reads the tokens; and the message of a failed parallel batch inside a caller's transaction does not claim that a +rollback removes rows the writers committed. Needs the native library: conftest.py collects this file on the native backend only. +""" + +import re +import time + +import numpy as np +import pytest + +from arcadedb._native import database as native_database +from arcadedb._native import extras +from arcadedb.exceptions import ArcadeDBError + +COLP = "arcadedb_n_insert_colp" + + +def _count(db, name): + sql = f"SELECT count(*) AS n FROM {name}" # nosec B608 - test-owned type name + return db.query("sql", sql).to_list()[0]["n"] + + +def _type(db, name="Col", unique=False): + db.command("sql", f"CREATE DOCUMENT TYPE {name} BUCKETS 3") + db.command("sql", f"CREATE PROPERTY {name}.id INTEGER") + db.command("sql", f"CREATE PROPERTY {name}.txt STRING") + if unique: + db.command("sql", f"CREATE INDEX ON {name} (id) UNIQUE") + + +class TokenSpy: + """The failure tokens that were made and the ones that were read (a read removes the token's entries in the image).""" + + def __init__(self, monkeypatch): + self.made, self.read = [], [] + real_make, real_read = extras.failure_token, extras.failures + real_rows = extras.insert_parallel + + def make(database): + fid = real_make(database) + self.made.append(fid) + return fid + + def failures(database, token): + self.read.append(token) + return real_read(database, token) + + def rows(database, type_name, chunk): + fid = real_rows(database, type_name, chunk) + self.made.append(fid["fid"]) + return fid + + monkeypatch.setattr(extras, "failure_token", make) + monkeypatch.setattr(extras, "failures", failures) + monkeypatch.setattr(extras, "insert_parallel", rows) + + +@pytest.fixture +def tokens(monkeypatch): + return TokenSpy(monkeypatch) + + +def _failed_count(error): + return int(re.search(r"reported (\d+) failed record", str(error)).group(1)) + + +def test_a_duplicate_key_is_reported_and_no_row_is_stored_twice(temp_db): + _type(temp_db, unique=True) + n, dupes = 3_000, 50 + ids = np.arange(n, dtype=np.int64) + ids[n - dupes :] = np.arange(dupes) # the last 50 repeat the first 50 keys + with pytest.raises(ArcadeDBError, match="failed record") as err: + temp_db.insert_columns("Col", {"id": ids, "txt": ["t"] * n}, parallel=True) + stored_ids = [r["id"] for r in temp_db.query("sql", "SELECT id FROM Col").to_list()] + stored = len(stored_ids) + assert ( + stored + _failed_count(err.value) >= n + ) # every row is a stored row or a reported one + # A writer that meets a duplicate key rolls back its WHOLE open batch (up to commit_every rows), so the valid rows of that batch are + # abandoned and reported too: `stored` can be below the number of distinct keys, never above it. + assert stored <= len(set(ids.tolist())) + assert set(stored_ids) <= set(ids.tolist()) # only ids that were sent + assert len(stored_ids) == len(set(stored_ids)) # none stored twice + + +def test_a_load_at_the_writers_batch_boundary_stores_every_row_and_nothing_later( + temp_db, +): + _type(temp_db) + threads = temp_db.async_executor().get_parallel_level() + n = threads * 10_240 + 1 # one row past a full batch for every writer + ids = np.arange(n, dtype=np.int64) + assert ( + temp_db.insert_columns("Col", {"id": ids, "txt": ["t"] * n}, parallel=True) == n + ) + stored = _count(temp_db, "Col") + assert stored == n + time.sleep(0.3) + temp_db.async_executor().wait_completion() + assert _count(temp_db, "Col") == n # nothing landed after the call returned + distinct = temp_db.query("sql", "SELECT count(DISTINCT id) AS c FROM Col") + assert distinct.to_list()[0]["c"] == n + + +def test_an_interrupt_in_the_crossing_waits_for_the_writers_and_reads_the_token( + temp_db, tokens, monkeypatch +): + _type(temp_db) + real_call = native_database._call + + def interrupted(prefix, name, *args): + reply = real_call(prefix, name, *args) + if name == COLP: + raise KeyboardInterrupt # the records are handed over; the reply never reaches the caller + return reply + + monkeypatch.setattr(native_database, "_call", interrupted) + n = 30_000 + with pytest.raises(KeyboardInterrupt): + temp_db.insert_columns( + "Col", {"id": np.arange(n, dtype=np.int64), "txt": ["t"] * n}, parallel=True + ) + assert ( + temp_db.async_executor().is_processing() is False + ) # waited: no writer runs on + assert len(tokens.made) == 1 and tokens.read == tokens.made # and the token is gone + assert ( + _count(temp_db, "Col") == n + ) # the rows handed over were all written, none half + + +def test_an_interrupt_between_the_chunks_of_the_row_path_waits_and_reads_the_tokens( + temp_db, tokens, monkeypatch +): + _type(temp_db) + real_rows = extras.insert_parallel + seen = [] + + def second_chunk_interrupted(database, type_name, chunk): + if seen: + raise KeyboardInterrupt + seen.append(1) + return real_rows(database, type_name, chunk) + + monkeypatch.setattr(extras, "insert_parallel", second_chunk_interrupted) + rows = [{"id": i, "txt": "t"} for i in range(25_000)] + with pytest.raises(KeyboardInterrupt): + temp_db.insert_many("Col", rows, parallel=True) + assert temp_db.async_executor().is_processing() is False + assert _count(temp_db, "Col") == 10_000 # the first chunk, whole + assert len(tokens.made) == 1 and tokens.read == tokens.made + + +def test_a_wait_that_fails_still_reads_the_tokens(temp_db, tokens, monkeypatch): + _type(temp_db) + real_wait = extras.NativeAsyncExecutor.wait_completion + state = {"raised": False} + + def failing_wait(self, timeout_ms=None): + if not state["raised"]: + state["raised"] = True + raise RuntimeError("the wait failed") + return real_wait(self, timeout_ms) + + monkeypatch.setattr(extras.NativeAsyncExecutor, "wait_completion", failing_wait) + n = 2_000 + with pytest.raises(ArcadeDBError, match="the wait failed"): + temp_db.insert_columns( + "Col", {"id": np.arange(n, dtype=np.int64), "txt": ["t"] * n}, parallel=True + ) + assert len(tokens.made) == 1 and tokens.read == tokens.made + temp_db.async_executor().wait_completion() + assert _count(temp_db, "Col") == n + + +def test_a_failed_parallel_batch_in_a_caller_transaction_does_not_promise_a_rollback( + temp_db, +): + import pyarrow as pa + + _type(temp_db, unique=True) + ids = list(range(300)) + ids[250] = 7 # a duplicate key in the third batch of 100 + table = pa.table({"id": pa.array(ids, pa.int64()), "txt": ["t"] * 300}) + temp_db.begin() + try: + with pytest.raises(ArcadeDBError) as err: + temp_db.insert_arrow("Col", table, batch_size=100, parallel=True) + message = str(err.value) + assert "not committed" not in message and "your transaction" in message + assert "async writers" in message and "does not remove them" in message + finally: + temp_db.rollback() + temp_db.async_executor().wait_completion() + assert ( + _count(temp_db, "Col") >= 200 + ) # the writers committed the first two batches: the rollback did not remove them