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..c703d872 100644 --- a/native/src/main/java/org/humemai/nativelib/BulkOps.java +++ b/native/src/main/java/org/humemai/nativelib/BulkOps.java @@ -10,16 +10,25 @@ 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.exception.NeedRetryException; +import com.arcadedb.schema.EdgeType; import com.arcadedb.schema.VertexType; import org.graalvm.nativeimage.c.type.CTypeConversion; import org.graalvm.word.Pointer; import org.graalvm.word.WordFactory; +import java.math.BigDecimal; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.ZoneOffset; import java.util.List; import java.util.Map; +import java.util.function.IntConsumer; final class BulkOps { private BulkOps() { @@ -55,26 +64,45 @@ 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) + "}"; } /** * header: {"n":rows,"commit":every,"cols":[[name,kind,dim],..]}. Kinds: q (int64), d (float64), b (bool byte), i (int32), e (float32), * f (float32 matrix, dim floats per row), s (UTF-8 string: int32 offsets[n+1] in the data slot, the bytes they index in the aux slot, as - * an Arrow string array; offsets[0] need not be 0). spec: per column six longs, [data address, data bytes, aux address, aux bytes, null + * an Arrow string array; offsets[0] need not be 0). The kinds that the JSON rows carry as tagged values (Json.decodeTag), read as the same + * objects: T (int64 ticks since 1970-01-01T00:00:00 of the UTC wall clock, dim = ticks per second: 1, 1000, 10^6 or 10^9, read as a + * LocalDateTime), D (int32 days since the epoch, a LocalDate), n (a decimal as text in the string layout, a BigDecimal) and x (bytes in + * the string layout, a byte[]), so a column of these is stored as the row path stores it. spec: per column six longs, [data address, data bytes, aux address, aux bytes, null * bitmap address, null bitmap bytes]; the bitmap is Arrow's (bit i of byte i/8, least significant first; 1 = present), an address of 0 * 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); @@ -102,13 +130,18 @@ private static long insertColumns(final Database db, final String type, final Ma final long need = switch (kinds[c]) { case 'q', 'd' -> 8L * n; case 'b' -> n; - case 'i', 'e' -> 4L * n; + case 'i', 'e', 'D' -> 4L * n; + case 'T' -> { + if (dims[c] != 1 && dims[c] != 1_000 && dims[c] != 1_000_000 && dims[c] != 1_000_000_000) + throw new IllegalArgumentException("column '" + names[c] + "': a date-time column has 1, 1000, 10^6 or 10^9 ticks per second, not " + dims[c]); + yield 8L * n; + } case 'f' -> { if (dims[c] <= 0) throw new IllegalArgumentException("column '" + names[c] + "': a vector column needs a dimension above 0"); yield 4L * n * dims[c]; } - case 's' -> 4L * (n + 1L); + case 's', 'n', 'x' -> 4L * (n + 1L); default -> throw new IllegalArgumentException("unknown column kind " + kinds[c]); }; if (dataBytes < need || (need > 0 && dataAddr == 0)) @@ -119,7 +152,7 @@ private static long insertColumns(final Database db, final String type, final Ma throw new IllegalArgumentException("column '" + names[c] + "' has a null bitmap of " + nullBytes + " bytes, " + (n + 7L) / 8 + " needed"); nulls[c] = nullAddr; } - if (kinds[c] == 's') { + if (kinds[c] == 's' || kinds[c] == 'n' || kinds[c] == 'x') { if (auxBytes < 0 || auxBytes > Integer.MAX_VALUE || (auxBytes > 0 && auxAddr == 0)) throw new IllegalArgumentException("column '" + names[c] + "': the string bytes are missing or beyond 2 GiB"); blobs[c] = auxBytes == 0 ? ByteBuffer.allocate(0) : CTypeConversion.asByteBuffer(WordFactory.pointer(auxAddr), (int) auxBytes); @@ -137,64 +170,157 @@ 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 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); - } - } + 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); } - doc.save(); - if (!wasActive && commitEvery > 0 && (i + 1) % commitEvery == 0) { - db.commit(); + } 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; + } + writeChunks(db, n, commitEvery, i -> { + final MutableDocument doc = vertex ? db.newVertex(type) : db.newDocument(type); + setRow(doc, i, names, kinds, dims, data, nulls, blobs, scratch); + doc.save(); + }); + return n; + } + + /** + * The retries of a chunk and the back-off before them are the engine's GraphBatch defaults (commitRetries 10, commitRetryDelayMs 1000, + * doubling up to 10 s), so a load that went through createVertices before keeps the same patience. + */ + private static final int COMMIT_RETRIES = 10; + private static final long COMMIT_RETRY_DELAY_MS = 1_000L; + private static final long MAX_COMMIT_RETRY_DELAY_MS = 10_000L; + + /** + * Writes rows [0, n) with writeRow (build and save row i), in transactions opened here: one per commitEvery rows (0: one for all). A + * chunk whose save or commit raises a NeedRetryException (an MVCC conflict with another writer, a lock timeout) is rolled back and + * written again from its first row after a back-off, up to COMMIT_RETRIES times; after the last attempt the error is thrown unchanged. + * Any other failure rolls the open chunk back and is thrown; the chunks committed before it stay. A transaction the caller opened is the + * caller's: the rows join it, nothing is committed, rolled back or retried here. + */ + private static void writeChunks(final Database db, final int n, final int commitEvery, final IntConsumer writeRow) { + if (db.isTransactionActive()) { + for (int i = 0; i < n; i++) + writeRow.accept(i); + return; + } + for (int from = 0; from < n;) { + final int to = commitEvery > 0 ? (int) Math.min((long) n, (long) from + commitEvery) : n; + for (int attempt = 0;;) { + try { db.begin(); + for (int i = from; i < to; i++) + writeRow.accept(i); + db.commit(); + break; + } catch (final NeedRetryException e) { + if (!rollbackQuietly(db, e) || ++attempt > COMMIT_RETRIES) + throw e; + backoffBeforeRetry(attempt); + } catch (final Throwable e) { + rollbackQuietly(db, e); + throw e; } } - if (!wasActive) - db.commit(); - } catch (final Throwable e) { - if (!wasActive && db.isTransactionActive()) { - try { - db.rollback(); - } catch (final Throwable rollbackError) { - e.addSuppressed(rollbackError); + from = to; + } + } + + /** Rolls back the open transaction, if any; false when the rollback itself failed (the failure is added to e as suppressed). */ + private static boolean rollbackQuietly(final Database db, final Throwable e) { + if (!db.isTransactionActive()) + return true; + try { + db.rollback(); + return true; + } catch (final Throwable rollbackError) { + e.addSuppressed(rollbackError); + return false; + } + } + + private static void backoffBeforeRetry(final int attempt) { + long delay = COMMIT_RETRY_DELAY_MS; + for (int i = 1; i < attempt && delay < MAX_COMMIT_RETRY_DELAY_MS; i++) + delay = Math.min(delay * 2, MAX_COMMIT_RETRY_DELAY_MS); + try { + Thread.sleep(delay); + } catch (final InterruptedException ie) { + Thread.currentThread().interrupt(); + throw new NeedRetryException("Interrupted while waiting to retry the commit of a bulk-insert chunk"); + } + } + + /** 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 'T' -> { + final long ticks = base.readLong(WordFactory.signed(8L * i)); + final long perSecond = dims[c]; + doc.set(names[c], LocalDateTime.ofEpochSecond(Math.floorDiv(ticks, perSecond), + (int) (Math.floorMod(ticks, perSecond) * (1_000_000_000L / perSecond)), ZoneOffset.UTC)); + } + case 'D' -> doc.set(names[c], LocalDate.ofEpochDay(base.readInt(WordFactory.signed(4L * i)))); + case 'n' -> { + final int from = base.readInt(WordFactory.signed(4L * i)); + final int len = base.readInt(WordFactory.signed(4L * (i + 1L))) - from; + blobs[c].get(from, scratch, 0, len); + doc.set(names[c], new BigDecimal(new String(scratch, 0, len, StandardCharsets.US_ASCII))); + } + case 'x' -> { + final int from = base.readInt(WordFactory.signed(4L * i)); + final byte[] bytes = new byte[base.readInt(WordFactory.signed(4L * (i + 1L))) - from]; + blobs[c].get(from, bytes, 0, bytes.length); + doc.set(names[c], bytes); + } + 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); } } - throw e; } - return n; } /** 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..c16d225b 100644 --- a/src/arcadedb/_native/bulk.py +++ b/src/arcadedb/_native/bulk.py @@ -1,12 +1,14 @@ """Columnar bulk insert through one crossing (``arcadedb_n_insert_colp``): whole columns go over by address, with no JSON text. Numpy 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. +vector columns, date-times, dates, decimals, bytes and Arrow arrays are laid out by :mod:`.columnar`; anything else (nested values, +mixed-type columns, 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; 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..d001b3da 100644 --- a/src/arcadedb/_native/columnar.py +++ b/src/arcadedb/_native/columnar.py @@ -2,19 +2,29 @@ A column crosses as raw memory, not as JSON text: the image reads each buffer in place and builds the documents without a parser. This module only decides how a column is laid out and hands the addresses over; it needs no native library, so the layouts are unit-tested -without one. A column that cannot be laid out (a nested value, a mixed-type column, a date-time, an integer past 64 bits, a string with -a lone surrogate) makes :func:`pack` return ``None`` and the caller uses the JSON path for that call; nothing is half sent. +without one. A column that cannot be laid out (a nested value, a mixed-type column, a numpy datetime64, an integer past 64 bits, a string +with a lone surrogate) makes :func:`pack` return ``None`` and the caller uses the JSON path for that call; nothing is half sent. Kinds (the image's, ``BulkOps.insertColumns``): ``q`` int64, ``d`` float64, ``b`` bool byte, ``i`` int32, ``e`` float32, ``f`` float32 matrix (``dim`` per row), ``s`` UTF-8 string (``int32 offsets[n+1]`` and the bytes they index: Arrow's string layout). A column may carry a null bitmap (Arrow's: bit ``i`` of byte ``i // 8``, least significant first, 1 = present). +Temporal, decimal and binary kinds (the values the JSON path sends as tagged values, so the stored value is the same object): ``T`` a +date-time, int64 ticks since 1970-01-01T00:00:00 of the UTC wall clock, ``dim`` ticks per second (1, 1000, 10**6 or 10**9: an Arrow +timestamp keeps its unit, a Python ``datetime`` is microseconds), read as ``LocalDateTime``; ``D`` a date, int32 days since the epoch, +read as ``LocalDate``; ``n`` a decimal, the string layout holding ``str(Decimal)``, read as ``BigDecimal``; ``x`` bytes, the string +layout, read as ``byte[]``. A time-zone-aware ``datetime`` or Arrow timestamp is its UTC wall clock and a naive one is taken as it is, +the rule of ``_tag_datetime``; nothing here decides a time zone. + ``i`` and ``e`` and the float matrix from a list of lists are chosen only for a property declared ``INTEGER``, ``FLOAT`` and ``ARRAY_OF_FLOATS``: the engine would convert to those types anyway, so the stored value is the one the JSON path stores. For an undeclared property an ``int32`` column crosses as ``q`` (a Long) and a ``float32`` column as ``d`` (a Double), as the JSON path stores them. """ +from datetime import date, datetime, timezone +from decimal import Decimal + from .schema_objects import ( # the schema-entry reader shared with NativeType (#195) _name_type, ) @@ -27,6 +37,16 @@ _INT64_MAX = 2**63 - 1 _INT32_MAX = 2**31 - 1 _NoneType = type(None) +_EPOCH_ORDINAL = date(1970, 1, 1).toordinal() +_UTC = timezone.utc +_BLOB_KINDS = ( + "snx" # the kinds laid out as int32 offsets[n + 1] plus the bytes they index +) +_TICKS = {"s": 1, "ms": 1_000, "us": 1_000_000, "ns": 1_000_000_000} +# the seconds that java.time.LocalDateTime.ofEpochSecond takes (years -999999999 to 999999999): beyond them it throws in the image, in the +# middle of a batch. Only a timestamp[s] can pass them: an int64 of milliseconds is at most 9.2e15 seconds +_LDT_MIN_SECONDS = -31_557_014_167_219_200 +_LDT_MAX_SECONDS = 31_556_889_864_403_199 class Packed: @@ -39,8 +59,8 @@ def __init__(self, name, kind, dim, array, aux=None, nulls=None, keep=()): self.kind = kind self.dim = dim self.n = int(array.shape[0]) - ( - 1 if kind == "s" else 0 - ) # a string column's array is its n + 1 offsets + 1 if kind in _BLOB_KINDS else 0 + ) # a string, decimal or bytes column's array is its n + 1 offsets self.data = (array.ctypes.data, array.nbytes) self.aux = aux or (0, 0) self.nulls = (nulls.ctypes.data, nulls.nbytes) if nulls is not None else (0, 0) @@ -166,6 +186,35 @@ def _string_column(name, texts, nulls=None): return Packed(name, "s", 0, offsets, aux=aux, nulls=nulls, keep=(blob_array, data)) +def _blob_column(name, kind, items, nulls=None): + """A list of bytes (a null already replaced by b'') as an offsets-plus-bytes column of ``kind``, or None past 2 GiB.""" + n = len(items) + data = b"".join(items) + if len(data) > _INT32_MAX: + return None + lengths = _np.fromiter(map(len, items), dtype=_np.int64, count=n) + offsets = _np.zeros(n + 1, dtype=_np.int32) + _np.cumsum(lengths, out=offsets[1:], dtype=_np.int32) + blob_array = ( + _np.frombuffer(data, dtype=_np.uint8) if data else _np.zeros(0, dtype=_np.uint8) + ) + aux = (blob_array.ctypes.data, blob_array.nbytes) if data else (0, 0) + return Packed(name, kind, 0, offsets, aux=aux, nulls=nulls, keep=(blob_array, data)) + + +def _epoch_micros(value): + """Microseconds since 1970-01-01T00:00:00 of the UTC wall clock of a ``datetime`` (a naive one is taken as it is, as ``_tag_datetime``).""" + if value.utcoffset() is not None: + value = value.astimezone(_UTC) + seconds = ( + (value.toordinal() - _EPOCH_ORDINAL) * 86400 + + value.hour * 3600 + + value.minute * 60 + + value.second + ) + return seconds * 1_000_000 + value.microsecond + + def _list_column(name, values, declared): """A list or tuple of Python values as a column: one type (None allowed) or a list of equal-length float arrays; else None.""" n = len(values) @@ -203,6 +252,33 @@ def _list_column(name, values, declared): return None # an integer beyond 64 bits kind = "q" if only is int else "d" if only is float else "b" return Packed(name, kind, 0, arr, nulls=nulls) + if issubclass(only, datetime): + if only.__name__ == "NaTType": + return None # the JSON path stores it as a null + try: + ticks = [0 if v is None else _epoch_micros(v) for v in values] + except (OverflowError, ValueError): + return None # the JSON path raises what it raised before + return Packed( + name, "T", 1_000_000, _np.array(ticks, dtype=_np.int64), nulls=nulls + ) + if only is date: + days = [0 if v is None else v.toordinal() - _EPOCH_ORDINAL for v in values] + return Packed(name, "D", 0, _np.array(days, dtype=_np.int32), nulls=nulls) + if only is Decimal: + texts = [] + for v in values: + if v is None: + texts.append(b"") + elif v.is_finite(): + texts.append(str(v).encode("ascii")) + else: + return None # the JSON path sends a non-finite Decimal down the per-row path + return _blob_column(name, "n", texts, nulls) + if only is bytes or only is bytearray: + return _blob_column( + name, "x", values if present is None else [v or b"" for v in values], nulls + ) if only is list and declared.get(name) == "ARRAY_OF_FLOATS" and present is None: try: arr = _np.array(values, dtype=_np.float32) @@ -248,10 +324,11 @@ def _arrow_column(name, arr, declared): if arr.null_count else None ) - if types.is_string(t) or types.is_large_string(t): + is_text = types.is_string(t) or types.is_large_string(t) + if is_text or types.is_binary(t) or types.is_large_binary(t): bufs = arr.buffers() off_buf, data_buf = bufs[1], bufs[2] - if types.is_string(t): + if types.is_string(t) or types.is_binary(t): offsets = _np.frombuffer(off_buf, dtype=_np.int32)[ arr.offset : arr.offset + n + 1 ] @@ -266,8 +343,42 @@ def _arrow_column(name, arr, declared): size = 0 if data_buf is None else min(data_buf.size, _INT32_MAX) aux = (data_buf.address, size) if size else (0, 0) return Packed( - name, "s", 0, offsets, aux=aux, nulls=nulls, keep=(arr, off_buf, data_buf) + name, + "s" if is_text else "x", + 0, + offsets, + aux=aux, + nulls=nulls, + keep=(arr, off_buf, data_buf), + ) + if types.is_timestamp(t): + values = arr.cast(pa.int64()).fill_null(0).to_numpy(zero_copy_only=False) + if t.unit == "s" and values.size: + if values.min() < _LDT_MIN_SECONDS or values.max() > _LDT_MAX_SECONDS: + return None # the row path raises its usual error for it, before anything is stored + return Packed( + name, + "T", + _TICKS[t.unit], + _np.ascontiguousarray(values, dtype=_np.int64), + nulls=nulls, + ) + if types.is_date(t): + if types.is_date64(t): # milliseconds since the epoch, floor-divided to days + raw = arr.cast(pa.int64()).fill_null(0).to_numpy(zero_copy_only=False) + days = raw // 86_400_000 + if days.size and ( + days.min() < -(_INT32_MAX + 1) or days.max() > _INT32_MAX + ): + return None # int32 days would wrap: the row path raises its usual error for it + else: + days = arr.cast(pa.int32()).fill_null(0).to_numpy(zero_copy_only=False) + return Packed( + name, "D", 0, _np.ascontiguousarray(days, dtype=_np.int32), nulls=nulls ) + if types.is_decimal(t): + texts = [b"" if v is None else str(v).encode("ascii") for v in arr.to_pylist()] + return _blob_column(name, "n", texts, nulls) if types.is_integer(t) or types.is_floating(t) or types.is_boolean(t): if arr.null_count: arr = arr.fill_null(False if types.is_boolean(t) else 0) @@ -343,17 +454,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..c52495df 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -877,7 +877,8 @@ def insert_many( a newDocument/set/save loop inside the image, as ``DocumentBatcher.insertManyJson`` on the JPype path). A ``datetime`` (also ``pandas.Timestamp`` and ``numpy.datetime64``), ``date``, ``Decimal``, ``bytes``/``bytearray``, numpy scalar, ``pandas.NA`` or ``pandas.NaT`` sitting directly in a row crosses in that text as a tagged value (``{"\\u0001": [kind, payload]}``, decoded by - ``Json.java``), so it does not move the chunk off this path; such a chunk loses the columnar entry. A row with a value the JSON + ``Json.java``), so it does not move the chunk off this path; but when a column holds only date-times, dates, Decimals or + bytes (None is a null) the chunk goes as columns, the same stored objects. A row with a value the JSON text cannot carry unchanged (NaN or Infinity, an integer beyond 64 bits, a date-time or decimal inside a list or map, a non-str key, a lone surrogate, a ``datetime.time``) goes through ``INSERT INTO t CONTENT :row`` on its own, in one crossing for a run of such rows; the rows around it keep the bulk path and every row is inserted in the order given. Inside the caller's @@ -995,20 +996,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 +1061,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 +1347,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 +1367,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) @@ -1351,7 +1394,7 @@ def insert_dataframe( ) -> int: """A pandas or Polars DataFrame into a type, in batches: the strict Arrow loading of ``Database.insert_dataframe`` (dtype refusals, declared types, NaN rules, the committed-rows message), batches through :meth:`insert_many` for a - document type and ``GraphBatch.create_vertices`` for a vertex type (``arcadedb/dataframe.py``). + document type and a vertex type alike (a vertex load refuses to run inside a transaction the caller opened; ``arcadedb/dataframe.py``). """ from ..dataframe import insert_table, to_arrow_table 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..6fffe536 100644 --- a/src/arcadedb/dataframe.py +++ b/src/arcadedb/dataframe.py @@ -5,17 +5,18 @@ (``Database.insert_dataframe`` and ``Database.insert_arrow``). Arrow is the one internal form: a pandas frame becomes an Arrow table, a Polars frame is already one, and each Arrow column becomes the typed array ``Database.insert_columns`` -sends across the bridge (a document type), or the property rows -``GraphBatch.create_vertices`` takes (a vertex type). No Java was added: both -paths use the existing glue (``DocumentBatcher.insertColumns`` and the -``GraphBatch`` vertex batcher). +sends across the bridge, for a document type and for a vertex type alike (the +vertex type gets vertices from the same ``DocumentBatcher.insertColumns``; it +refuses to run inside a transaction the caller opened, as ``GraphBatch`` does). +No Java was added. The native backend (the ``arcadedb`` wheel) shares everything but the last step: the checks, the type mapping, the batching and the error wording are -the same code. A document batch goes over as one JSON text through -``Database.insert_many`` (one crossing per 10,000 rows), a vertex batch through -``GraphBatch.create_vertices``, and a value JSON cannot carry (a date-time, a -decimal, bytes) travels as a tagged value the image decodes. +the same code. A batch of documents or vertices goes over by address through +the columnar entry (``arcadedb_n_insert_colp``) when every column can cross as +memory, else as one JSON text through ``Database.insert_many`` (one crossing per +10,000 rows), where a value JSON cannot carry (a date-time, a decimal, bytes) +travels as a tagged value the image decodes. """ import re @@ -321,12 +322,28 @@ def _scalars_to_tagged(kind: str, chunk): ] +def _is_number_list(pa, arrow_type) -> bool: + """A list, large list or fixed-size list of integers or floats: the column the columnar entry lays out as a float32 matrix.""" + types = pa.types + if not ( + types.is_list(arrow_type) + or types.is_large_list(arrow_type) + or types.is_fixed_size_list(arrow_type) + ): + return False + return types.is_floating(arrow_type.value_type) or types.is_integer( + arrow_type.value_type + ) + + def _column_values(name: str, chunk, pa, native: bool = False, arrow: bool = False): """One Arrow array as what ``Database.insert_columns`` takes: a numpy array when the column crosses as one buffer (integers, floats and booleans without nulls), else a list of values (``None`` is a null). - ``arrow`` (the native backend's document load): a string column, and a numeric one with nulls, stay Arrow arrays, which - the columnar entry reads in place (the offsets, the bytes and the validity bitmap) instead of as Python values. + ``arrow`` (the native backend's document and vertex load): a string, binary, timestamp, date or decimal column, a list of + numbers (a vector) and a numeric column with nulls, stay Arrow arrays, which the columnar entry reads in place (the offsets, + the bytes, the integer ticks, the vector values and the validity bitmap) instead of as Python values; ``_insert_documents_native`` + turns them into Python or tagged values for the row path when a batch cannot go columnar. """ types = pa.types if types.is_dictionary(chunk.type): # a pandas categorical or a Polars Categorical @@ -361,6 +378,10 @@ def _column_values(name: str, chunk, pa, native: bool = False, arrow: bool = Fal values = values.astype("float64") # float16/32 widen exactly return values return chunk.to_pylist() + if arrow and kind in ("timestamp", "date", "decimal", "binary"): + return chunk + if arrow and kind == "list" and _is_number_list(pa, chunk.type): + return chunk # a vector column: the columnar entry reads the values in place (to_pylist makes one Python float per element) if kind in ("timestamp", "date"): return ( _temporal_to_tagged(chunk, pa) if native else _temporal_to_java(chunk, pa) @@ -421,68 +442,79 @@ 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()) - graph_batch = db.graph_batch() if is_vertex else None - try: - for batch in table.to_batches(max_chunksize=batch_size): - if batch.num_rows == 0: - continue - columns = { - name: _column_values( - name, - batch.column(i), - pa, - native, - arrow=native and not is_vertex and not parallel, + # nothing then); a vertex load refuses to run inside one, as GraphBatch.createVertices does. + # 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_transaction = bool(db.is_transaction_active()) + in_caller_transaction = in_transaction and not parallel + if is_vertex and in_transaction: + raise ArcadeDBError( + f"Failed to create batch vertices for type '{type_name}': a vertex load commits its own batches " + f"and cannot run inside a transaction the caller opened: commit or roll back that transaction first" + ) + for batch in table.to_batches(max_chunksize=batch_size): + if batch.num_rows == 0: + continue + columns = { + name: _column_values( + name, + batch.column(i), + pa, + native, + arrow=native, + ) + for i, name in enumerate(names) + } + try: + if native: + _insert_documents_native( + db, type_name, names, columns, batch_size, parallel, declared + ) + else: + db.insert_columns( + type_name, columns, commit_every=batch_size, parallel=parallel ) - for i, name in enumerate(names) - } - try: - if is_vertex: - _insert_vertices( - graph_batch, type_name, names, columns, batch.num_rows - ) - elif native: - _insert_documents_native( - db, type_name, names, columns, batch_size, parallel, declared - ) - else: - db.insert_columns( - type_name, columns, commit_every=batch_size, parallel=parallel - ) - except ArcadeDBError as exc: - if written and in_caller_transaction: - # the batches joined the caller's transaction: nothing is committed - # until the caller commits, and a rollback removes them - raise ArcadeDBError( - f"{exc}; the {written} row(s) of the batches before this one " - f"were written in your transaction, not committed" - ) from exc - if written: - raise ArcadeDBError( - f"{exc}; the {written} row(s) of the batches before this one " - f"were committed" - ) from exc - raise - written += batch.num_rows - finally: - if graph_batch is not None: - graph_batch.close() + except ArcadeDBError as exc: + if written and in_caller_transaction: + # the batches joined the caller's transaction: nothing is committed + # until the caller commits, and a rollback removes them + raise ArcadeDBError( + 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 " + f"were committed" + ) from exc + raise + written += batch.num_rows + del ( + columns, + batch, + ) # sent: drop this batch's columns before the next batch is built, so two batches are never alive together return written -def _insert_vertices( - graph_batch, type_name: str, names: List[str], columns: Dict[str, Any], n: int -): - """One batch of vertices: the columns as one property dict per row, through ``GraphBatch.create_vertices``.""" - lists = [] - for name in names: - values = columns[name] - lists.append(values.tolist() if hasattr(values, "tolist") else values) - graph_batch.create_vertices( - type_name, [dict(zip(names, row)) for row in zip(*lists)] - ) +def _arrow_row_values(col) -> list: + """A column kept as an Arrow array, as the values the row path takes: Python values, or tagged values for a timestamp, date, + decimal or binary column (the form it had before the columnar entry took them, so the stored value is unchanged). + """ + import pyarrow as pa + + kind = _kind(pa, col.type) + if kind in ("timestamp", "date"): + return _temporal_to_tagged(col, pa) + if kind in ("decimal", "binary"): + return _scalars_to_tagged(kind, col) + return col.to_pylist() def _insert_documents_native( @@ -497,9 +529,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 ( @@ -520,7 +563,7 @@ def _insert_documents_native( if hasattr( col, "to_pylist" ): # a column kept as an Arrow array for the columnar entry - lists.append(col.to_pylist()) + lists.append(_arrow_row_values(col)) else: # jvm-payload: not-a-crossing -- no JVM here: insert_many serialises the rows to JSON text lists.append(col.tolist() if hasattr(col, "tolist") else col) diff --git a/src/java/com/arcadedb/python/DocumentBatcher.java b/src/java/com/arcadedb/python/DocumentBatcher.java index bbcae4f8..278361b5 100644 --- a/src/java/com/arcadedb/python/DocumentBatcher.java +++ b/src/java/com/arcadedb/python/DocumentBatcher.java @@ -31,6 +31,7 @@ import com.arcadedb.database.Database; import com.arcadedb.database.MutableDocument; import com.arcadedb.database.async.ErrorCallback; +import com.arcadedb.exception.NeedRetryException; import com.arcadedb.schema.DocumentType; import com.arcadedb.schema.EdgeType; import com.arcadedb.schema.VertexType; @@ -38,6 +39,7 @@ import com.arcadedb.serializer.json.JSONObject; import java.util.Map; +import java.util.function.IntConsumer; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; @@ -69,31 +71,11 @@ public static long insertManyJson(final Database db, final String typeName, fina // commits and the failure rollback apply only to the transactions opened // here, matching the Python per-row fallback (#7882). final boolean vertex = n > 0 && isVertexType(db, typeName); - final boolean wasActive = db.isTransactionActive(); - if (!wasActive) - db.begin(); - try { - for (int i = 0; i < n; i++) { - final MutableDocument doc = newRecord(db, typeName, vertex); - fill(doc, rows.getJSONObject(i), i, vectorNames, vectors); - doc.save(); - if (!wasActive && commitEvery > 0 && (i + 1) % commitEvery == 0) { - db.commit(); - db.begin(); - } - } - if (!wasActive) - db.commit(); - } catch (final Throwable e) { - if (!wasActive && db.isTransactionActive()) { - try { - db.rollback(); - } catch (final Throwable rollbackError) { - e.addSuppressed(rollbackError); - } - } - throw e; - } + writeChunks(db, n, commitEvery, i -> { + final MutableDocument doc = newRecord(db, typeName, vertex); + fill(doc, rows.getJSONObject(i), i, vectorNames, vectors); + doc.save(); + }); return n; } @@ -113,32 +95,80 @@ public static long insertColumns(final Database db, final String typeName, final final int n, final int commitEvery) { checkColumns(names, columns, n); final boolean vertex = n > 0 && isVertexType(db, typeName); - final boolean wasActive = db.isTransactionActive(); - if (!wasActive) - db.begin(); - try { - for (int i = 0; i < n; i++) { - final MutableDocument doc = newRecord(db, typeName, vertex); - fillRow(doc, names, columns, i); - doc.save(); - if (!wasActive && commitEvery > 0 && (i + 1) % commitEvery == 0) { - db.commit(); - db.begin(); - } - } - if (!wasActive) - db.commit(); - } catch (final Throwable e) { - if (!wasActive && db.isTransactionActive()) { + writeChunks(db, n, commitEvery, i -> { + final MutableDocument doc = newRecord(db, typeName, vertex); + fillRow(doc, names, columns, i); + doc.save(); + }); + return n; + } + + /** + * The retries of a chunk and the back-off before them are the engine's GraphBatch defaults (commitRetries 10, commitRetryDelayMs 1000, + * doubling up to 10 s), so a load that went through createVertices before keeps the same patience. + */ + private static final int COMMIT_RETRIES = 10; + private static final long COMMIT_RETRY_DELAY_MS = 1_000L; + private static final long MAX_COMMIT_RETRY_DELAY_MS = 10_000L; + + /** + * Writes rows [0, n) with writeRow (build and save row i), in transactions opened here: one per commitEvery rows (0: one for all). A + * chunk whose save or commit raises a NeedRetryException (an MVCC conflict with another writer, a lock timeout) is rolled back and + * written again from its first row after a back-off, up to COMMIT_RETRIES times; after the last attempt the error is thrown unchanged. + * Any other failure rolls the open chunk back and is thrown; the chunks committed before it stay. A transaction the caller opened is the + * caller's: the rows join it, nothing is committed, rolled back or retried here. + */ + private static void writeChunks(final Database db, final int n, final int commitEvery, final IntConsumer writeRow) { + if (db.isTransactionActive()) { + for (int i = 0; i < n; i++) + writeRow.accept(i); + return; + } + for (int from = 0; from < n;) { + final int to = commitEvery > 0 ? (int) Math.min((long) n, (long) from + commitEvery) : n; + for (int attempt = 0;;) { try { - db.rollback(); - } catch (final Throwable rollbackError) { - e.addSuppressed(rollbackError); + db.begin(); + for (int i = from; i < to; i++) + writeRow.accept(i); + db.commit(); + break; + } catch (final NeedRetryException e) { + if (!rollbackQuietly(db, e) || ++attempt > COMMIT_RETRIES) + throw e; + backoffBeforeRetry(attempt); + } catch (final Throwable e) { + rollbackQuietly(db, e); + throw e; } } - throw e; + from = to; + } + } + + /** Rolls back the open transaction, if any; false when the rollback itself failed (the failure is added to e as suppressed). */ + private static boolean rollbackQuietly(final Database db, final Throwable e) { + if (!db.isTransactionActive()) + return true; + try { + db.rollback(); + return true; + } catch (final Throwable rollbackError) { + e.addSuppressed(rollbackError); + return false; + } + } + + private static void backoffBeforeRetry(final int attempt) { + long delay = COMMIT_RETRY_DELAY_MS; + for (int i = 1; i < attempt && delay < MAX_COMMIT_RETRY_DELAY_MS; i++) + delay = Math.min(delay * 2, MAX_COMMIT_RETRY_DELAY_MS); + try { + Thread.sleep(delay); + } catch (final InterruptedException ie) { + Thread.currentThread().interrupt(); + throw new NeedRetryException("Interrupted while waiting to retry the commit of a bulk-insert chunk"); } - return n; } /** diff --git a/tests/conftest.py b/tests/conftest.py index ec3147c7..af18bb6f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -47,6 +47,18 @@ 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") + # a column of date-times, dates, Decimals or bytes crosses as a column (kinds T D n x), observed on the native library + collect_ignore.append("test_native_columnar_temporal.py") + # the review of the temporal kinds: a value the image cannot take is refused in Python, never wrapped or thrown mid-batch + collect_ignore.append("test_native_columnar_temporal_review_fixes.py") + # a vertex load by insert_dataframe or insert_arrow crosses as columns, one crossing per batch, observed on the native library + collect_ignore.append("test_native_vertex_columnar.py") + # a document or vertex load keeps its transient Python memory to one batch and a vector column is never turned into Python floats + collect_ignore.append("test_native_columnar_memory.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_columnar_temporal_lists.py b/tests/test_columnar_temporal_lists.py new file mode 100644 index 00000000..9a9d62b7 --- /dev/null +++ b/tests/test_columnar_temporal_lists.py @@ -0,0 +1,114 @@ +"""A column of date-times, dates, Decimals or bytes stores the same values on both wheels (ArcadeData/arcadedb-python #316). + +On the native wheel such a column used to move the whole batch off the columnar entry; now ``insert_columns`` and ``insert_many`` send +it as a column (``tests/test_native_columnar_temporal.py`` proves the entry and compares with the row path). Here the expected values +are literals, so the same file on the JVM wheel and on the native wheel says that both store them: a naive date-time as it is, a +time-zone-aware one as its UTC wall clock, a date-time before 1970, a null, a date, a Decimal and bytes. +""" + +import datetime as dt +import decimal + +import pytest + +UTC9 = dt.timezone(dt.timedelta(hours=9)) +TS = [ + dt.datetime(2020, 1, 2, 3, 4, 5, 123456), + dt.datetime(1969, 12, 31, 23, 59, 59, 999999), + dt.datetime(1900, 1, 1, 0, 0, 0, 1), + dt.datetime(2400, 2, 29, 23, 59, 59), + None, +] +TS_STORED = TS # a naive date-time is the wall clock +AWARE = [ + dt.datetime(2020, 1, 2, 12, 0, 0, 5, tzinfo=UTC9), + dt.datetime(1969, 12, 31, 23, 0, 0, tzinfo=dt.timezone.utc), + dt.datetime(2020, 6, 1, 0, 0, 0, tzinfo=dt.timezone(dt.timedelta(hours=-7))), + None, + dt.datetime(2000, 1, 1, tzinfo=dt.timezone.utc), +] +AWARE_STORED = [ + dt.datetime(2020, 1, 2, 3, 0, 0, 5), + dt.datetime(1969, 12, 31, 23, 0, 0), + dt.datetime(2020, 6, 1, 7, 0, 0), + None, + dt.datetime(2000, 1, 1), +] +DAYS = [ + dt.date(2020, 2, 29), + dt.date(1960, 1, 1), + None, + dt.date(1970, 1, 1), + dt.date(1969, 12, 31), +] +AMOUNTS = [ + decimal.Decimal("12345.6789"), + decimal.Decimal("-0.50"), + decimal.Decimal("1E+3"), + None, + decimal.Decimal("0"), +] +BLOBS = [b"\x00\xff abc\x7f", b"", None, bytearray(b"ab\x80"), b"z"] + + +def _type(db, name): + db.command("sql", f"CREATE DOCUMENT TYPE {name}") + db.command("sql", f"CREATE PROPERTY {name}.id INTEGER") + db.command("sql", f"CREATE PROPERTY {name}.ts DATETIME_MICROS") + db.command("sql", f"CREATE PROPERTY {name}.day DATE") + db.command("sql", f"CREATE PROPERTY {name}.amount DECIMAL") + db.command("sql", f"CREATE PROPERTY {name}.blob BINARY") + + +def _stored(db, name, column): + sql = f"SELECT id, {column} AS v FROM {name} ORDER BY id" # nosec B608 - test-owned names + return [r.get("v") for r in db.query("sql", sql).to_list()] + + +def _bytes(values): + # a BINARY property reads back as the values of a Java byte[] (signed, -128..127) or as bytes; & 0xFF takes both to the same bytes + return [None if v is None else bytes(b & 0xFF for b in v) for v in values] + + +@pytest.mark.parametrize("path", ["insert_columns", "insert_many"]) +def test_naive_and_aware_datetimes_dates_decimals_and_bytes_are_stored_the_same( + temp_db, path +): + _type(temp_db, "Ev") + n = len(TS) + columns = { + "id": list(range(n)), + "ts": TS, + "day": DAYS, + "amount": AMOUNTS, + "blob": BLOBS, + } + if path == "insert_columns": + assert temp_db.insert_columns("Ev", columns) == n + else: + rows = [dict(zip(columns, vals)) for vals in zip(*columns.values())] + assert temp_db.insert_many("Ev", rows) == n + assert _stored(temp_db, "Ev", "ts") == TS_STORED + assert _stored(temp_db, "Ev", "day") == DAYS + assert _stored(temp_db, "Ev", "amount") == AMOUNTS + assert _bytes(_stored(temp_db, "Ev", "blob")) == _bytes(BLOBS) + + +def test_time_zone_aware_datetimes_are_stored_as_their_utc_wall_clock(temp_db): + _type(temp_db, "Zoned") + n = len(AWARE) + assert temp_db.insert_columns("Zoned", {"id": list(range(n)), "ts": AWARE}) == n + assert _stored(temp_db, "Zoned", "ts") == AWARE_STORED + + +def test_a_parallel_load_of_datetimes_stores_the_same_values(temp_db): + _type(temp_db, "Par") + n = len(TS) + assert ( + temp_db.insert_columns( + "Par", {"id": list(range(n)), "ts": TS, "day": DAYS}, parallel=True + ) + == n + ) + assert _stored(temp_db, "Par", "ts") == TS_STORED + assert _stored(temp_db, "Par", "day") == DAYS diff --git a/tests/test_native_columnar_insert.py b/tests/test_native_columnar_insert.py index 1ec2984c..e7296960 100644 --- a/tests/test_native_columnar_insert.py +++ b/tests/test_native_columnar_insert.py @@ -315,15 +315,18 @@ def test_rows_with_different_keys_fall_back_to_json(temp_db, spy, monkeypatch): _fallback_case(temp_db, spy, monkeypatch, rows, ROWS) -def test_a_datetime_leaves_the_columnar_entry_but_stays_on_the_json_entry( - temp_db, spy, monkeypatch -): - """columnar.pack refuses a date-time, so the chunk takes the JSON entry, which carries it as a tagged value - (test_native_mixed_chunk.py compares what it stores with INSERT ... CONTENT).""" +def test_a_datetime_column_stays_on_the_columnar_entry(temp_db, spy, monkeypatch): + """A column of date-times is the kind T of the columnar entry (#316): the stored value is the one the JSON entry stores as a tagged + value (tests/test_native_columnar_temporal.py has the full comparison, test_native_mixed_chunk.py the rows that stay on JSON). + """ when = datetime.datetime(2024, 5, 6, 7, 8, 9) rows = [{"id": i, "txt": "t%d" % i, "when": when} for i in range(10)] - _fallback_case(temp_db, spy, monkeypatch, rows, ROWS) - assert spy.count(MANY) == 0 + _types(temp_db, "Ref", "Col") + _json_reference(temp_db, "Ref", rows, monkeypatch) + spy.calls.clear() + temp_db.insert_many("Col", rows) + assert spy.count(COLP) == 1 and spy.count(ROWS) == 0 and spy.count(MANY) == 0 + assert _stored(temp_db, "Col") == _stored(temp_db, "Ref") def test_values_json_cannot_carry_still_go_row_by_row_in_one_crossing( @@ -344,9 +347,7 @@ def test_an_integer_past_64_bits_is_refused_on_both_paths(temp_db, spy): temp_db.insert_columns("Col", {"id": np.array([2**63], dtype=np.uint64)}) -def test_a_dataframe_with_a_column_the_buffers_cannot_carry_falls_back( - temp_db, spy, monkeypatch -): +def test_a_dataframe_with_a_timestamp_column_stays_columnar(temp_db, spy, monkeypatch): import pandas as pd import pyarrow as pa @@ -362,7 +363,7 @@ def test_a_dataframe_with_a_column_the_buffers_cannot_carry_falls_back( } ) temp_db.insert_arrow("Col", table) - assert spy.count(COLP) == 0 + assert spy.count(COLP) == 1 and spy.count(ROWS) == 0 assert pd is not None rows = temp_db.query("sql", "SELECT id, txt FROM Col ORDER BY id").to_list() assert [(r["id"], r["txt"]) for r in rows] == [(1, "a"), (2, "b"), (3, "c")] diff --git a/tests/test_native_columnar_memory.py b/tests/test_native_columnar_memory.py new file mode 100644 index 00000000..b3ed404c --- /dev/null +++ b/tests/test_native_columnar_memory.py @@ -0,0 +1,300 @@ +"""Native: a document or vertex load by ``insert_arrow`` or ``insert_dataframe`` keeps its transient memory to one batch, and a vector column +is never turned into Python floats on its way to the columnar entry (ArcadeData/arcadedb-python #333 and #335, the extra anonymous memory of +the post1 candidate's native wheel). + +Three properties are pinned, none of them with an absolute size: + +* ``_column_values(..., arrow=True)`` returns an Arrow list of numbers (a vector column) as the Arrow array, which ``columnar.pack`` reads in + place; the JVM build (``arrow=False``) still gets Python lists. One ``to_pylist()`` of a 64-dimensional vector column makes 64 Python + floats per row (about 2 KB a row), which a 10,000-row batch held for the whole crossing. +* 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 + objects, not Arrow's memory or the image's heap, which is why the ratio is generous). +* ``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 + is sent (the arrays of batch k are gone before batch k + 1 is packed). + +Needs the native library, so conftest.py collects this file on the native backend only. +""" + +import ctypes +import gc +import tracemalloc +import weakref + +import numpy as np +import pyarrow as pa +import pytest + +from arcadedb import dataframe +from arcadedb._native import columnar + +DIM = 64 + + +def _vectors(n, start=0, dim=DIM): + r = np.arange(start, start + n, dtype=np.int64) + j = np.arange(dim, dtype=np.int64) + return ( + ((r[:, None] * 131 + j[None, :] * 17) % 1000).astype(np.float32) + ) / np.float32(1000.0) + + +def _fixed(mat): + return pa.FixedSizeListArray.from_arrays( + pa.array(mat.reshape(-1), type=pa.float32()), mat.shape[1] + ) + + +def _bytes(addr_len): + addr, length = addr_len + return ctypes.string_at(addr, length) if length else b"" + + +def _peak(fn): + gc.collect() + tracemalloc.start() + try: + tracemalloc.reset_peak() + base = tracemalloc.get_traced_memory()[0] + fn() + return tracemalloc.get_traced_memory()[1] - base + finally: + tracemalloc.stop() + + +@pytest.mark.parametrize( + "make", + [ + lambda m: _fixed(m), + lambda m: pa.array(m.astype(np.float64).tolist(), type=pa.list_(pa.float64())), + lambda m: pa.array(m.astype(np.int32).tolist(), type=pa.large_list(pa.int32())), + ], + ids=["fixed_size_list_f32", "list_f64", "large_list_i32"], +) +def test_a_vector_column_stays_an_arrow_array_on_the_native_load(make): + chunk = make(_vectors(10)) + kept = dataframe._column_values("embedding", chunk, pa, native=True, arrow=True) + assert isinstance(kept, pa.Array) and kept.type == chunk.type + assert ( + dataframe._column_values("embedding", chunk, pa, native=True, arrow=False) + == chunk.to_pylist() + ) + assert isinstance( + dataframe._column_values("embedding", chunk, pa, native=False, arrow=False), + list, + ) + + +def test_a_list_that_is_not_a_list_of_numbers_still_becomes_python_values(): + for chunk in ( + pa.array([["a", "b"], ["c"]]), + pa.array([[{"x": 1}], [{"x": 2}]]), + pa.array([[[1.0], [2.0]], [[3.0]]]), + ): + got = dataframe._column_values("c", chunk, pa, native=True, arrow=True) + assert isinstance(got, list) and got == chunk.to_pylist() + + +def test_the_vector_layout_from_arrow_is_the_layout_from_python_lists(): + mat = _vectors(300) + table = pa.table({"embedding": _fixed(mat)}).slice( + 7, 211 + ) # a slice: the offset of the array is not zero + declared = columnar.Declared(known={"embedding": "ARRAY_OF_FLOATS"}) + kept = dataframe._column_values( + "embedding", table.column(0).combine_chunks(), pa, True, True + ) + listed = dataframe._column_values( + "embedding", table.column(0).combine_chunks(), pa, True, False + ) + rows_a, (col_a,) = columnar.pack({"embedding": kept}, declared) + rows_b, (col_b,) = columnar.pack({"embedding": listed}, declared) + assert ( + rows_a == rows_b == 211 + and col_a.kind == col_b.kind == "f" + and col_a.dim == col_b.dim == DIM + ) + assert _bytes(col_a.data) == _bytes(col_b.data) == mat[7:218].tobytes() + + +def test_the_python_side_peak_of_a_vector_batch_is_a_fraction_of_the_python_float_copy(): + batch = _fixed(_vectors(10_000)) + declared = columnar.Declared(known={"embedding": "ARRAY_OF_FLOATS"}) + as_python = _peak( + lambda: columnar.pack( + { + "embedding": dataframe._column_values( + "embedding", batch, pa, True, False + ) + }, + declared, + ) + ) + in_place = _peak( + lambda: columnar.pack( + {"embedding": dataframe._column_values("embedding", batch, pa, True, True)}, + declared, + ) + ) + assert ( + as_python > 10 * 1024 * 1024 + ), as_python # 640,000 Python floats: the copy this change removes (about 20 MiB) + assert in_place * 8 < as_python, ( + in_place, + as_python, + ) # measured: about 100 times less + + +class _PackSpy: + """Wraps ``columnar.pack``: the row count of every call and a weak reference to every array it laid out.""" + + def __init__(self, monkeypatch): + self.rows, self.arrays, self.alive_at_next_call = [], [], [] + real = columnar.pack + + def pack(columns, declared): + gc.collect() + self.alive_at_next_call.append( + sum(ref() is not None for refs in self.arrays for ref in refs) + ) + out = real(columns, declared) + self.rows.append(None if out is None else out[0]) + self.arrays.append( + [weakref.ref(c.keep[0]) for c in out[1]] if out is not None else [] + ) + return out + + monkeypatch.setattr(columnar, "pack", pack) + + +def _vector_type(db, kind="VERTEX"): + db.command("sql", f"CREATE {kind} TYPE V") + db.command("sql", "CREATE PROPERTY V.id LONG") + db.command("sql", "CREATE PROPERTY V.embedding ARRAY_OF_FLOATS") + + +@pytest.mark.parametrize("kind", ["VERTEX", "DOCUMENT"]) +def test_every_batch_is_packed_once_and_released_before_the_next( + temp_db, monkeypatch, kind +): + _vector_type(temp_db, kind) + spy = _PackSpy(monkeypatch) + n, batch = 2_505, 1_000 + table = pa.table( + {"id": pa.array(np.arange(n, dtype=np.int64)), "embedding": _fixed(_vectors(n))} + ) + assert temp_db.insert_arrow("V", table, batch_size=batch) == n + assert spy.rows == [ + 1_000, + 1_000, + 505, + ] # one pack per batch, never more than batch_size rows + assert spy.alive_at_next_call == [ + 0, + 0, + 0, + ], spy.alive_at_next_call # batch k is gone when batch k + 1 is packed + gc.collect() + assert not any( + ref() is not None for refs in spy.arrays for ref in refs + ) # and the last one when the call returns + + +def test_a_dataframe_load_is_packed_per_batch_too(temp_db, monkeypatch): + import pandas as pd + + temp_db.command("sql", "CREATE VERTEX TYPE P") + temp_db.command("sql", "CREATE PROPERTY P.id LONG") + temp_db.command("sql", "CREATE PROPERTY P.name STRING") + spy = _PackSpy(monkeypatch) + frame = pd.DataFrame( + {"id": np.arange(250, dtype=np.int64), "name": [f"p{i}" for i in range(250)]} + ) + assert temp_db.insert_dataframe("P", frame, batch_size=100) == 250 + assert spy.rows == [100, 100, 50] and spy.alive_at_next_call == [0, 0, 0] + + +def test_the_python_side_peak_of_a_whole_vector_load_stays_below_one_batch_of_python_floats( + temp_db, +): + _vector_type(temp_db) + n, batch = 40_000, 10_000 + table = pa.table( + {"id": pa.array(np.arange(n, dtype=np.int64)), "embedding": _fixed(_vectors(n))} + ) + one_batch_of_python_floats = _peak( + lambda: table.slice(0, batch).column("embedding").to_pylist() + ) + whole_load = _peak(lambda: temp_db.insert_arrow("V", table, batch_size=batch)) + assert one_batch_of_python_floats > 10 * 1024 * 1024 + assert whole_load * 2 < one_batch_of_python_floats, ( + whole_load, + one_batch_of_python_floats, + ) # measured: about 7 times less + assert temp_db.query("sql", "SELECT count(*) AS n FROM V").first().get("n") == n + + +def _stored(db, key): + return [ + float(x) + for x in db.query("sql", "SELECT embedding FROM V WHERE id = ?", key) + .first() + .get("embedding") + ] + + +@pytest.mark.parametrize("kind", ["VERTEX", "DOCUMENT"]) +def test_the_stored_vectors_are_the_ones_the_python_lists_stored(temp_db, kind): + _vector_type(temp_db, kind) + mat = _vectors(120) + table = pa.table( + {"id": pa.array(np.arange(120, dtype=np.int64)), "embedding": _fixed(mat)} + ) + assert ( + temp_db.insert_arrow("V", table.slice(5, 100), batch_size=30) == 100 + ) # a sliced table, four batches + for key in (5, 6, 34, 35, 64, 104): + assert ( + max(abs(a - b) for a, b in zip(_stored(temp_db, key), mat[key].tolist())) + < 1e-6 + ) + assert temp_db.query("sql", "SELECT count(*) AS n FROM V").first().get("n") == 100 + + +def test_a_vector_column_with_a_null_or_a_list_property_takes_the_row_path_with_the_same_values( + temp_db, +): + _vector_type(temp_db) + temp_db.command("sql", "CREATE VERTEX TYPE L") + temp_db.command("sql", "CREATE PROPERTY L.id LONG") + temp_db.command("sql", "CREATE PROPERTY L.tags LIST") + with_null = pa.table( + { + "id": pa.array([1, 2, 3], type=pa.int64()), + "embedding": pa.array( + [[0.5, 0.25], None, [1.0, 2.0]], type=pa.list_(pa.float32()) + ), + } + ) + assert temp_db.insert_arrow("V", with_null) == 3 + assert [ + float(x) + for x in temp_db.query("sql", "SELECT embedding FROM V WHERE id = 1") + .first() + .get("embedding") + ] == [0.5, 0.25] + assert ( + temp_db.query("sql", "SELECT embedding FROM V WHERE id = 2") + .first() + .get("embedding") + is None + ) + tags = pa.table( + { + "id": pa.array([1, 2], type=pa.int64()), + "tags": pa.array([[1, 2, 3], [4]], type=pa.list_(pa.int64())), + } + ) + assert temp_db.insert_arrow("L", tags) == 2 + assert list( + temp_db.query("sql", "SELECT tags FROM L WHERE id = 1").first().get("tags") + ) == [1, 2, 3] diff --git a/tests/test_native_columnar_pack.py b/tests/test_native_columnar_pack.py index cda0a1f7..e8925661 100644 --- a/tests/test_native_columnar_pack.py +++ b/tests/test_native_columnar_pack.py @@ -190,7 +190,11 @@ def test_a_bitmap_longer_than_a_byte(): [1, 2**63], # past 64 bits [{"a": 1}, {"a": 2}], # nested [[1], [2]], # a list without a declared ARRAY_OF_FLOATS - [b"ab", b"cd"], + [ + b"ab", + "cd", + ], # bytes and text (a column of bytes alone crosses as kind x since #316) + [b"ab", bytearray(b"cd")], # two bytes-like types [np.int64(1), np.int64(2)], # numpy scalars in a list are not Python numbers ], ) @@ -198,16 +202,107 @@ def test_columns_the_buffer_cannot_carry_unchanged_are_refused(column): assert columnar.pack({"c": column}, NO_TYPES) is None -def test_datetimes_decimals_and_object_arrays_of_them_are_refused(): +def test_numpy_datetime64_arrays_and_mixed_temporal_columns_are_refused(): import datetime import decimal - assert columnar.pack({"d": [datetime.datetime(2020, 1, 1)]}, NO_TYPES) is None - assert columnar.pack({"d": [decimal.Decimal("1.5")]}, NO_TYPES) is None assert ( columnar.pack({"d": np.array(["2020-01-01"], dtype="datetime64[D]")}, NO_TYPES) is None ) + assert ( + columnar.pack( + {"d": [datetime.datetime(2020, 1, 1), datetime.date(2020, 1, 1)]}, NO_TYPES + ) + is None + ) + assert ( + columnar.pack({"d": [decimal.Decimal("1"), decimal.Decimal("NaN")]}, NO_TYPES) + is None + ) + assert columnar.pack({"d": [b"a", "b"]}, NO_TYPES) is None + + +def _ticks(col, dtype): + return np.frombuffer(_bytes(col.data), dtype=dtype).tolist() + + +def test_python_datetimes_are_utc_wall_clock_microseconds(): + import datetime + + utc9 = datetime.timezone(datetime.timedelta(hours=9)) + values = [ + datetime.datetime(1970, 1, 1), + datetime.datetime(1969, 12, 31, 23, 59, 59, 999999), + datetime.datetime(2020, 1, 2, 12, 0, 0, 5, tzinfo=utc9), # 03:00:00.000005 UTC + None, + ] + rows, cols = columnar.pack({"t": values}, NO_TYPES) + col = cols[0] + assert rows == 4 and col.kind == "T" and col.dim == 1_000_000 + assert _ticks(col, np.int64) == [0, -1, 1_577_934_000_000_005, 0] + assert _present(col) == [True, True, True, False] + + +def test_python_dates_are_epoch_days(): + import datetime + + rows, cols = columnar.pack( + {"d": [datetime.date(1970, 1, 2), datetime.date(1969, 12, 31), None]}, NO_TYPES + ) + assert cols[0].kind == "D" and _ticks(cols[0], np.int32) == [1, -1, 0] + assert _present(cols[0]) == [True, True, False] + + +def test_decimals_and_bytes_use_the_string_layout(): + import decimal + + rows, cols = columnar.pack( + { + "a": [decimal.Decimal("1E+3"), None, decimal.Decimal("-0.50")], + "b": [b"\x00\xff", b"", None], + "c": [bytearray(b"xy"), bytearray(b"z"), bytearray(b"")], + }, + NO_TYPES, + ) + a, b, c = cols + assert [x.kind for x in cols] == ["n", "x", "x"] + offsets = np.frombuffer(_bytes(a.data), dtype=np.int32).tolist() + assert offsets == [0, 4, 4, 9] and _bytes(a.aux) == b"1E+3-0.50" + offsets = np.frombuffer(_bytes(b.data), dtype=np.int32).tolist() + assert offsets == [0, 2, 2, 2] and _bytes(b.aux) == b"\x00\xff" + assert _bytes(c.aux) == b"xyz" and _present(a) == [True, False, True] + + +def test_arrow_timestamps_dates_decimals_and_binary(): + import decimal + + import pyarrow as pa + + stamps = pa.array([1, None, -2], pa.int64()).cast( + pa.timestamp("ms", tz="Asia/Seoul") + ) + days = pa.array([3, -4, None], pa.int32()).cast(pa.date32()) + day_ms = pa.array([86_400_000 * 2, -1, None], pa.int64()).cast(pa.date64()) + amounts = pa.array([decimal.Decimal("1.50"), None], pa.decimal128(10, 2)) + blobs = pa.array([b"ab", None, b"c"]) + for name, col, expected in ( + ("t", stamps, ("T", 1_000)), + ("d", days, ("D", 0)), + ("m", day_ms, ("D", 0)), + ): + packed = columnar.pack({name: col}, NO_TYPES)[1][0] + assert (packed.kind, packed.dim) == expected + assert _ticks(columnar.pack({"t": stamps}, NO_TYPES)[1][0], np.int64) == [1, 0, -2] + assert _ticks(columnar.pack({"d": days}, NO_TYPES)[1][0], np.int32) == [3, -4, 0] + # a date64 is floor-divided to days: 2 days, then -1 ms is the day before the epoch + assert _ticks(columnar.pack({"m": day_ms}, NO_TYPES)[1][0], np.int32) == [2, -1, 0] + n = columnar.pack({"n": amounts}, NO_TYPES)[1][0] + assert n.kind == "n" and _bytes(n.aux) == b"1.50" + x = columnar.pack({"x": blobs}, NO_TYPES)[1][0] + assert ( + x.kind == "x" and _bytes(x.aux) == b"abc" and _present(x) == [True, False, True] + ) def test_ragged_columns_and_empty_input_are_left_to_the_row_path(): @@ -233,6 +328,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 diff --git a/tests/test_native_columnar_temporal.py b/tests/test_native_columnar_temporal.py new file mode 100644 index 00000000..7ae3b545 --- /dev/null +++ b/tests/test_native_columnar_temporal.py @@ -0,0 +1,326 @@ +"""Native: a column of date-times, dates, Decimals or bytes stays on the columnar entry and stores what the row path stores (#316). + +Before, one such column made ``columnar.pack`` return None, so every column of the batch went through one dict per row and a JSON text +with tagged values. Now the kinds T, D, n and x of ``arcadedb_n_insert_colp`` carry them and the image builds the same objects +(``LocalDateTime``, ``LocalDate``, ``BigDecimal``, ``byte[]``) that ``Json.decodeTag`` builds for the tagged values. Each test loads the +same data through the row path (the private reference: the columnar path switched off, so the chunk goes as tagged JSON) and through +the columns, compares every stored value and its type, and spies on the entry points so a silent fall back cannot pass. Needs the +native library, so conftest.py collects this file on the native backend only. +""" + +import datetime as dt +import decimal + +import pytest + +from arcadedb._native import bulk, columnar, lib # noqa: E402 + +COLP, ROWS, MANY = ( + "arcadedb_n_insert_colp", + "arcadedb_n_insert_rows", + "arcadedb_n_exec_many", +) +UTC9 = dt.timezone(dt.timedelta(hours=9)) + + +class Spy: + def __init__(self, monkeypatch): + self.calls = [] + real = lib.call + + def call(name, *args): + self.calls.append(name) + return real(name, *args) + + monkeypatch.setattr(lib, "call", call) + + def count(self, name): + return self.calls.count(name) + + +@pytest.fixture +def spy(monkeypatch): + return Spy(monkeypatch) + + +def _type(db, name, declared=True): + db.command("sql", f"CREATE DOCUMENT TYPE {name}") + db.command("sql", f"CREATE PROPERTY {name}.id INTEGER") + if declared: + for prop, ty in ( + ("ts", "DATETIME_MICROS"), + ("ns", "DATETIME_NANOS"), + ("day", "DATE"), + ("amount", "DECIMAL"), + ("blob", "BINARY"), + ): + db.command("sql", f"CREATE PROPERTY {name}.{prop} {ty}") + + +def _stored(db, name): + """Every record by id as {property: (python type name, repr)}, so a changed type is as visible as a changed value.""" + sql = f"SELECT FROM {name} ORDER BY id" # nosec B608 - test-owned type name + return [ + { + k: (type(v).__name__, repr(v)) + for k, v in sorted(rec.items()) + if not k.startswith("@") and k != "rid" + } + for rec in db.query("sql", sql).to_list() + ] + + +def _rows(columns): + return [dict(zip(columns, vals)) for vals in zip(*columns.values())] + + +def _reference(db, name, rows, monkeypatch): + with monkeypatch.context() as m: + m.setattr(bulk, "_ENABLED", False) + db.insert_many(name, rows) + + +def _columns(n=60): + base = dt.datetime(1969, 12, 31, 23, 59, 59, 999999) + return { + "id": list(range(n)), + "ts": [ + ( + None + if i % 11 == 0 + else (base + dt.timedelta(seconds=i * 86_399, microseconds=i * 7)) + ) + for i in range(n) + ], + "day": [ + None if i % 13 == 0 else dt.date(1960, 1, 1) + dt.timedelta(days=i * 401) + for i in range(n) + ], + "amount": [ + ( + None + if i % 9 == 0 + else decimal.Decimal(f"{i * 1234567}.{i % 100:02d}") + * (-1 if i % 2 else 1) + ) + for i in range(n) + ], + "blob": [ + None if i % 8 == 0 else bytes(range(i % 256)) + b"\x00\xff" + for i in range(n) + ], + } + + +@pytest.mark.parametrize("declared", [True, False], ids=["declared", "schemaless"]) +def test_a_batch_with_temporal_decimal_and_bytes_columns_stays_columnar( + temp_db, spy, monkeypatch, declared +): + for name in ("Ref", "ByRows", "ByColumns"): + _type(temp_db, name, declared) + columns = _columns() + _reference(temp_db, "Ref", _rows(columns), monkeypatch) + assert ( + spy.count(COLP) == 0 and spy.count(ROWS) == 1 + ) # the reference is the JSON path + spy.calls.clear() + assert temp_db.insert_many("ByRows", _rows(columns)) == 60 + assert spy.count(COLP) == 1 and spy.count(ROWS) == 0 and spy.count(MANY) == 0 + spy.calls.clear() + assert temp_db.insert_columns("ByColumns", columns) == 60 + assert spy.count(COLP) == 1 and spy.count(ROWS) == 0 + expected = _stored(temp_db, "Ref") + assert len(expected) == 60 and expected[1]["ts"][0] == "datetime" + assert _stored(temp_db, "ByRows") == expected + assert _stored(temp_db, "ByColumns") == expected + + +def test_time_zone_aware_and_naive_datetimes_follow_the_json_rule( + temp_db, spy, monkeypatch +): + """An aware date-time is its UTC wall clock, a naive one is taken as it is (_tag_datetime): the column does not decide a zone.""" + _type(temp_db, "Ref") + _type(temp_db, "Col") + values = [ + dt.datetime(2020, 1, 2, 12, 0, 0, 5, tzinfo=UTC9), + dt.datetime(1969, 12, 31, 23, 0, 0, tzinfo=dt.timezone.utc), + dt.datetime(2020, 6, 1, 0, 0, 0, tzinfo=dt.timezone(dt.timedelta(hours=-7))), + ] + rows = [{"id": i, "ts": v} for i, v in enumerate(values)] + _reference(temp_db, "Ref", rows, monkeypatch) + spy.calls.clear() + temp_db.insert_many("Col", rows) + assert spy.count(COLP) == 1 + assert _stored(temp_db, "Col") == _stored(temp_db, "Ref") + + +def test_pandas_timestamps_keep_microseconds_as_the_json_path_does( + temp_db, spy, monkeypatch +): + import pandas as pd + + _type(temp_db, "Ref") + _type(temp_db, "Col") + stamps = [ + pd.Timestamp("2026-05-06 07:08:09.123456789"), + pd.Timestamp("1960-05-06 07:08:09.123456", tz="Asia/Seoul"), + ] + rows = [{"id": i, "ts": v} for i, v in enumerate(stamps)] + _reference(temp_db, "Ref", rows, monkeypatch) + spy.calls.clear() + temp_db.insert_many("Col", rows) + assert spy.count(COLP) == 1 and spy.count(ROWS) == 0 + assert _stored(temp_db, "Col") == _stored(temp_db, "Ref") + + +@pytest.mark.parametrize( + "column, colp_runs", + [ + ([dt.datetime(2020, 1, 1), dt.date(2020, 1, 1)], 0), # a date-time and a date + # A non-finite Decimal: the column is refused (BigDecimal("NaN") would throw mid-batch in the image), so the NaN row goes to the + # statement on its own, as before. The row around it is a run of carriable rows: it goes in bulk, and the columnar entry + # carries a finite Decimal now, so that one run crosses on arcadedb_n_insert_colp (before, it was the JSON entry). + ([decimal.Decimal("1.5"), decimal.Decimal("NaN")], 1), + ([b"a", "b"], 0), # bytes and text + ], + ids=["datetime-and-date", "nan-decimal", "bytes-and-text"], +) +def test_a_column_the_buffers_cannot_carry_unchanged_still_falls_back( + temp_db, spy, monkeypatch, column, colp_runs +): + assert columnar.pack({"x": column}, columnar.Declared(known={})) is None + _type(temp_db, "Ref", declared=False) + _type(temp_db, "Col", declared=False) + rows = [{"id": i, "x": v} for i, v in enumerate(column)] + try: + _reference(temp_db, "Ref", rows, monkeypatch) + except ( + Exception + ): # the statement path may refuse it; then the new path must refuse it the same way + with pytest.raises(Exception): + temp_db.insert_many("Col", rows) + assert spy.count(COLP) == colp_runs + return + spy.calls.clear() + temp_db.insert_many("Col", rows) + assert spy.count(COLP) == colp_runs + assert _stored(temp_db, "Col") == _stored(temp_db, "Ref") + + +def _arrow_table(pa): + import numpy as np + + n = 40 + ids = np.arange(n, dtype=np.int64) + epoch_us = (ids - 20) * 1_000_003_000 + 7 # microseconds, before and after 1970 + mask = ids % 6 == 0 + + def stamps(unit, tz=None): + raw = { + "s": epoch_us // 1_000_000, + "ms": epoch_us // 1_000, + "us": epoch_us, + "ns": epoch_us * 1_000, + }[unit] + return pa.array(raw.astype(np.int64), mask=mask).cast(pa.timestamp(unit, tz=tz)) + + return pa.table( + { + "id": ids, + "ts": stamps("us"), + "ns": stamps("ns"), + "sec": stamps("s"), + "zoned": stamps("ms", tz="Asia/Seoul"), + "day": pa.array( + ((ids - 10) * 397).astype(np.int32), mask=ids % 5 == 0 + ).cast(pa.date32()), + "amount": pa.array( + [ + None if i % 7 == 0 else decimal.Decimal(f"{i}.{i % 100:02d}") + for i in range(n) + ], + type=pa.decimal128(20, 2), + ), + "blob": pa.array( + [None if i % 4 == 0 else bytes(range(i)) for i in range(n)] + ), + } + ) + + +def test_arrow_timestamps_dates_decimals_and_binary_stay_columnar( + temp_db, spy, monkeypatch +): + import pyarrow as pa + + for name in ("Ref", "Col"): + _type(temp_db, name) + for prop, ty in (("sec", "DATETIME_SECOND"), ("zoned", "DATETIME")): + temp_db.command("sql", f"CREATE PROPERTY {name}.{prop} {ty}") + table = _arrow_table(pa) + with ( + monkeypatch.context() as m + ): # the reference: tagged values through the JSON path + m.setattr(bulk, "_ENABLED", False) + assert temp_db.insert_arrow("Ref", table) == 40 + assert spy.count(COLP) == 0 + spy.calls.clear() + assert temp_db.insert_arrow("Col", table) == 40 + assert spy.count(COLP) == 1 and spy.count(ROWS) == 0 and spy.count(MANY) == 0 + expected = _stored(temp_db, "Ref") + assert len(expected) == 40 and expected[1]["ns"][0] == "datetime" + assert _stored(temp_db, "Col") == expected + + +def test_an_arrow_batch_with_one_list_column_falls_back_with_the_same_values( + temp_db, spy, monkeypatch +): + """A list column cannot cross as a buffer: the whole batch takes the row path, and the temporal columns reach it as tagged values.""" + import pyarrow as pa + + for name in ("Ref", "Col"): + _type(temp_db, name) + for prop, ty in (("sec", "DATETIME_SECOND"), ("zoned", "DATETIME")): + temp_db.command("sql", f"CREATE PROPERTY {name}.{prop} {ty}") + table = _arrow_table(pa).append_column( + "tags", pa.array([["a", "b"]] * 40, type=pa.list_(pa.string())) + ) + with monkeypatch.context() as m: + m.setattr(bulk, "_ENABLED", False) + temp_db.insert_arrow("Ref", table) + spy.calls.clear() + temp_db.insert_arrow("Col", table) + assert spy.count(COLP) == 0 + assert _stored(temp_db, "Col") == _stored(temp_db, "Ref") + + +def test_a_parallel_load_of_datetimes_is_one_columnar_crossing(temp_db, spy): + _type(temp_db, "Par") + columns = _columns(300) + assert temp_db.insert_columns("Par", columns, parallel=True) == 300 + assert spy.count(COLP) == 1 and spy.count(ROWS) == 0 + assert temp_db.query("sql", "SELECT count(*) AS n FROM Par").first().get("n") == 300 + + +def test_the_image_refuses_a_wrong_ticks_per_second_and_writes_nothing(temp_db): + import json + + _type(temp_db, "Col") + rows, packed = columnar.pack( + {"id": [1, 2], "ts": [dt.datetime(2020, 1, 1)] * 2}, columnar.Declared(known={}) + ) + assert packed[1].kind == "T" and packed[1].dim == 1_000_000 + packed[1].dim = 7 + spec = columnar.spec_of(packed) + reply = json.loads( + lib.call( + COLP, + temp_db._h, + lib.b("Col"), + lib.b(columnar.header_of(rows, 10_000, packed)), + spec, + len(spec), + ) + ) + assert not reply["ok"] and "ticks per second" in reply["error"] + assert temp_db.query("sql", "SELECT count(*) AS n FROM Col").first().get("n") == 0 diff --git a/tests/test_native_columnar_temporal_review_fixes.py b/tests/test_native_columnar_temporal_review_fixes.py new file mode 100644 index 00000000..34898f9c --- /dev/null +++ b/tests/test_native_columnar_temporal_review_fixes.py @@ -0,0 +1,190 @@ +"""Review of the temporal columnar kinds (#333): a value the image cannot take is refused in Python, before the crossing, and never wraps. + +``T`` and ``D`` columns carry integers that the image turns into ``LocalDateTime`` and ``LocalDate``. A ``date64`` past int32 days used +to wrap in the int32 cast and a ``timestamp[s]`` past the years of ``java.time`` threw ``DateTimeException`` in the middle of a batch. Now +``columnar.pack`` returns None for both, so the call takes the row path, which raises what it raised before (nothing stored, no +transaction left open on the synchronous call; ``may have been stored`` for a parallel call, which has handed rows over by then). Also +pinned: a chunk with a date-time column and a float column holding NaN goes columnar now and stores the same NaN that ``INSERT INTO t +CONTENT :row`` stores (both end in ``MutableDocument.putProperty``). Needs the native library: conftest.py collects this file on the +native backend only. +""" + +import datetime as dt +import decimal +import math + +import numpy as np +import pytest + +from arcadedb._native import columnar, lib # noqa: E402 +from arcadedb.exceptions import ArcadeDBError + +pa = pytest.importorskip("pyarrow") + +COLP = "arcadedb_n_insert_colp" +DAY_MS = 86_400_000 +LDT_MIN_SECONDS, LDT_MAX_SECONDS = -31_557_014_167_219_200, 31_556_889_864_403_199 +NO_TYPES = columnar.Declared(known={}) + + +class Spy: + def __init__(self, monkeypatch): + self.calls = [] + real = lib.call + + def call(name, *args): + self.calls.append(name) + return real(name, *args) + + monkeypatch.setattr(lib, "call", call) + + +def _count(db, name): + sql = f"SELECT count(*) AS n FROM {name}" # nosec B608 - test-owned type name + return db.query("sql", sql).first().get("n") + + +def _kinds(columns): + packed = columnar.pack(columns, NO_TYPES) + return None if packed is None else [c.kind for c in packed[1]] + + +def test_a_date64_past_int32_days_is_not_laid_out_and_never_wraps(): + wraps = pa.array([DAY_MS * (2**31 + 5)], pa.int64()).cast(pa.date64()) + assert columnar.pack({"d": wraps}, NO_TYPES) is None + below = pa.array([DAY_MS * (-(2**31) - 1)], pa.int64()).cast(pa.date64()) + assert columnar.pack({"d": below}, NO_TYPES) is None + one_bad = pa.array([0, DAY_MS, DAY_MS * 2**33], pa.int64()).cast(pa.date64()) + assert columnar.pack({"d": one_bad}, NO_TYPES) is None # one value is enough + + +def test_a_date64_at_the_int32_limits_stays_columnar_and_exact(): + edge = pa.array( + [DAY_MS * (2**31 - 1), DAY_MS * -(2**31), None], pa.int64() + ).cast(pa.date64()) + rows, packed = columnar.pack({"d": edge}, NO_TYPES) + assert rows == 3 and packed[0].kind == "D" + import ctypes + + import numpy as np + + addr, length = packed[0].data + days = np.frombuffer(ctypes.string_at(addr, length), dtype=np.int32) + assert days.tolist() == [2**31 - 1, -(2**31), 0] + + +def test_a_timestamp_past_the_years_of_java_time_is_not_laid_out(): + assert ( + columnar.pack({"t": pa.array([0, 2**62], pa.timestamp("s"))}, NO_TYPES) is None + ) + over = pa.array([LDT_MAX_SECONDS + 1], pa.timestamp("s")) + assert columnar.pack({"t": over}, NO_TYPES) is None + under = pa.array([LDT_MIN_SECONDS - 1], pa.timestamp("s")) + assert columnar.pack({"t": under}, NO_TYPES) is None + + +def test_a_timestamp_at_the_limits_stays_columnar(): + edge = pa.array([LDT_MAX_SECONDS, LDT_MIN_SECONDS, None], pa.timestamp("s")) + assert _kinds({"t": edge}) == ["T"] + # an int64 of milliseconds, microseconds or nanoseconds is always inside the years of java.time + for unit in ("ms", "us", "ns"): + big = pa.array([2**63 - 1, -(2**63)], pa.timestamp(unit)) + assert _kinds({"t": big}) == ["T"], unit + + +@pytest.mark.parametrize("text", ["NaN", "-NaN", "sNaN", "Infinity", "-Infinity"]) +def test_a_non_finite_decimal_is_never_laid_out(text): + # new BigDecimal("NaN") throws in the image, in the middle of a batch: the column must stay in Python, in every container + bad = decimal.Decimal(text) + assert columnar.pack({"x": [decimal.Decimal("1.5"), bad]}, NO_TYPES) is None + assert columnar.pack({"x": [None, bad, None]}, NO_TYPES) is None + as_object = np.array([decimal.Decimal("1"), bad], dtype=object) + assert columnar.pack({"x": as_object}, NO_TYPES) is None + finite = columnar.pack({"x": [decimal.Decimal("1.5"), None]}, NO_TYPES) + assert finite is not None and finite[1][0].kind == "n" # the finite column crosses + + +def _type(db, name): + db.command("sql", f"CREATE DOCUMENT TYPE {name}") + db.command("sql", f"CREATE PROPERTY {name}.id INTEGER") + + +def test_a_timestamp_the_image_cannot_take_stores_nothing_and_leaves_no_transaction( + temp_db, +): + _type(temp_db, "Col") + table = pa.table( + { + "id": pa.array([1, 2], pa.int64()), + "t": pa.array([0, 2**62], pa.timestamp("s")), + } + ) + with pytest.raises(ArcadeDBError): + temp_db.insert_arrow("Col", table) + assert not temp_db.is_transaction_active() + assert _count(temp_db, "Col") == 0 + + +def test_a_date_the_image_cannot_take_stores_nothing_and_leaves_no_transaction(temp_db): + _type(temp_db, "Col") + wraps = pa.array([0, DAY_MS * (2**31 + 5)], pa.int64()).cast(pa.date64()) + table = pa.table({"id": pa.array([1, 2], pa.int64()), "d": wraps}) + with pytest.raises(ArcadeDBError): + temp_db.insert_arrow("Col", table) + assert not temp_db.is_transaction_active() + assert _count(temp_db, "Col") == 0 + + +def test_a_parallel_load_of_such_a_timestamp_says_that_rows_may_have_been_stored( + temp_db, +): + _type(temp_db, "Col") + table = pa.table( + { + "id": pa.array([1, 2], pa.int64()), + "t": pa.array([0, 2**62], pa.timestamp("s")), + } + ) + with pytest.raises(ArcadeDBError, match="may have been stored"): + temp_db.insert_arrow("Col", table, parallel=True) + temp_db.async_executor().wait_completion() + assert ( + _count(temp_db, "Col") <= 1 + ) # the row before the bad one may have been handed over + + +def test_a_datetime_chunk_with_nan_stores_the_nan_the_statement_path_stores( + temp_db, monkeypatch +): + for name in ("Col", "Ref"): + _type(temp_db, name) + temp_db.command("sql", f"CREATE PROPERTY {name}.ts DATETIME_MICROS") + temp_db.command("sql", f"CREATE PROPERTY {name}.val DOUBLE") + rows = [ + { + "id": i, + "ts": dt.datetime(2020, 1, 1, 0, 0, i), + "val": float("nan") if i % 2 else 1.5 * i, + } + for i in range(10) + ] + spy = Spy(monkeypatch) + temp_db.insert_many("Col", rows) + assert ( + spy.calls.count(COLP) == 1 + ) # the date-time column no longer forces the JSON path + sql = "INSERT INTO `Ref` CONTENT :row" + temp_db.begin() + temp_db._command_many("sql", sql, [{"row": row} for row in rows]) + temp_db.commit() + + def stored(name): + select = f"SELECT FROM {name} ORDER BY id" # nosec B608 - test-owned type name + return [ + (r["id"], type(r["val"]).__name__, repr(r["val"]), repr(r["ts"])) + for r in temp_db.query("sql", select).to_list() + ] + + got, want = stored("Col"), stored("Ref") + assert got == want + assert all(math.isnan(float(v[2])) for v in got if v[0] % 2) diff --git a/tests/test_native_mixed_chunk.py b/tests/test_native_mixed_chunk.py index db4162cf..cf409cb6 100644 --- a/tests/test_native_mixed_chunk.py +++ b/tests/test_native_mixed_chunk.py @@ -169,7 +169,8 @@ def test_a_pandas_timestamp_is_stored_as_the_statement_path_stores_it( _statement_path(db, "Ref", rows) spy.clear() db.insert_many("Bulk", rows) - assert spy.count(MANY) == 0 and spy.count(ROWS) == 1 + # one column of Timestamps is the kind T of the columnar entry (#316): microseconds, as the tagged path keeps them + assert spy.count(MANY) == 0 and spy.count(ROWS) == 0 and spy.count(COLP) == 1 assert _stored(db, "Bulk") == _stored(db, "Ref") @@ -249,8 +250,10 @@ def test_an_uncarriable_value_sends_its_own_row_to_the_statement_not_the_chunk( spy.clear() db.insert_many("Bulk", rows) assert spy.count(MANY) == 1, "one statement run for the one uncarriable row" + # the rows before and after it stay on the bulk entry; its columnar form carries the date-time column now, so each run + # crosses on arcadedb_n_insert_colp (before it was the JSON entry, arcadedb_n_insert_rows) assert ( - spy.count(ROWS) == 2 + spy.count(COLP) == 2 and spy.count(ROWS) == 0 ), "the rows before and after it stay on the bulk entry" assert _stored(db, "Bulk") == _stored(db, "Ref") diff --git a/tests/test_native_vertex_columnar.py b/tests/test_native_vertex_columnar.py new file mode 100644 index 00000000..154a1fe2 --- /dev/null +++ b/tests/test_native_vertex_columnar.py @@ -0,0 +1,55 @@ +"""Native: a vertex type loaded by ``insert_dataframe`` or ``insert_arrow`` crosses as columns (``arcadedb_n_insert_colp``), one crossing +per batch, with no GraphBatch crossing and no JSON text (ArcadeData/arcadedb-python #319). The stored vertices are compared with the old +route in ``tests/test_vertex_load_columnar.py``, which runs on both wheels. Needs the native library, so conftest.py collects this file on +the native backend only. +""" + +import numpy as np +import pandas as pd +import pyarrow as pa + +from arcadedb._native import lib # noqa: E402 + +COLP, ROWS, MANY = ( + "arcadedb_n_insert_colp", + "arcadedb_n_insert_rows", + "arcadedb_n_exec_many", +) + + +def _spy(monkeypatch): + calls = [] + real = lib.call + + def call(name, *args): + calls.append(name) + return real(name, *args) + + monkeypatch.setattr(lib, "call", call) + return calls + + +def _vertex_type(db, name): + db.command("sql", f"CREATE VERTEX TYPE {name}") + db.command("sql", f"CREATE PROPERTY {name}.id LONG") + + +def test_every_batch_of_a_vertex_load_is_one_columnar_crossing(temp_db, monkeypatch): + _vertex_type(temp_db, "V") + calls = _spy(monkeypatch) + frame = pd.DataFrame( + { + "id": np.arange(25, dtype=np.int64), + "name": [f"p{i}" for i in range(25)], + "score": [None if i % 5 == 0 else i * 0.5 for i in range(25)], + "born": pd.to_datetime(["2000-01-01"] * 25), + } + ) + assert temp_db.insert_dataframe("V", frame, batch_size=10) == 25 + assert calls.count(COLP) == 3 # batches of 10, 10 and 5 + assert calls.count(ROWS) == 0 and calls.count(MANY) == 0 + assert not [c for c in calls if "gb" in c or "graph" in c], calls + assert temp_db.query("sql", "SELECT count(*) AS n FROM V").first().get("n") == 25 + calls.clear() + assert temp_db.insert_arrow("V", pa.Table.from_pandas(frame)) == 25 + assert calls.count(COLP) == 1 diff --git a/tests/test_vertex_load_columnar.py b/tests/test_vertex_load_columnar.py new file mode 100644 index 00000000..e5a8d87c --- /dev/null +++ b/tests/test_vertex_load_columnar.py @@ -0,0 +1,178 @@ +"""``insert_dataframe`` and ``insert_arrow`` load a vertex type through the columnar insert, and the stored vertices, their indexes and +their buckets are those of the old route, ``GraphBatch.create_vertices`` (ArcadeData/arcadedb-python #319). + +Both routes are ``newVertex``, ``set`` per property and ``save`` on the caller's thread inside a transaction, so the records, the +index entries and the bucket each vertex lands in come out the same; the one difference the engine code shows is that a GraphBatch +turns the write-ahead log off for its load (its ``useWAL`` default) while the columnar insert uses the database's own settings, as the +document loads of the same calls always did. The old route is the private reference here: a type loaded through +``graph_batch().create_vertices`` and a type loaded through ``insert_dataframe`` are compared value by value and type by type, with the +index lookups and the per-bucket record counts. A spy on ``create_vertices`` proves which route ran (the old code calls it, the new code +must not). Runs on both wheels. +""" + +import datetime as dt + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest + +from arcadedb.exceptions import ArcadeDBError + +N = 37 # not a multiple of the batch size or of the bucket count + + +def _type(db, name, buckets=3): + db.command("sql", f"CREATE VERTEX TYPE {name} BUCKETS {buckets}") + db.command("sql", f"CREATE PROPERTY {name}.id LONG") + db.command("sql", f"CREATE PROPERTY {name}.grp INTEGER") + db.command("sql", f"CREATE PROPERTY {name}.emb ARRAY_OF_FLOATS") + db.command("sql", f"CREATE INDEX ON {name} (id) UNIQUE") + db.command("sql", f"CREATE INDEX ON {name} (grp) NOTUNIQUE") + + +def _frame(n=N): + ids = np.arange(n, dtype=np.int64) + return pd.DataFrame( + { + "id": ids, + "grp": (ids % 5).astype(np.int32), + "name": [None if i % 9 == 0 else f"n{i}-é" for i in range(n)], + "score": [None if i % 4 == 0 else i * 0.25 for i in range(n)], + "born": pd.to_datetime( + [ + dt.datetime(1969, 12, 31, 23, 59, 59, 999999) + dt.timedelta(days=i) + for i in range(n) + ] + ), + "emb": [np.array([i, i + 0.5, -i], dtype=np.float32) for i in range(n)], + } + ) + + +def _rows_of(frame): + out = [] + for rec in frame.to_dict("records"): + row = {} + for k, v in rec.items(): + if v is None or (isinstance(v, float) and v != v): + v = None # an Arrow null, as the old route of insert_dataframe handed it over + elif isinstance(v, pd.Timestamp): + v = v.to_pydatetime() + elif isinstance(v, np.ndarray): + v = v.tolist() + elif isinstance(v, np.generic): + v = v.item() + row[k] = v + out.append(row) + return out + + +def _old_route(db, name, frame): + with db.graph_batch() as batch: + batch.create_vertices(name, _rows_of(frame)) + + +def _stored(db, name): + sql = f"SELECT FROM {name} ORDER BY id" # nosec B608 - test-owned type name + return [ + {k: (type(v).__name__, v) for k, v in rec.items() if not k.startswith("@")} + for rec in db.query("sql", sql).to_list() + ] + + +def _bucket_counts(db, name): + counts = [] + for i in range(3): + sql = f"SELECT count(*) AS n FROM bucket:{name}_{i}" # nosec B608 - test-owned names + counts.append(db.query("sql", sql).first().get("n")) + return sorted(counts) + + +def _index_view(db, name): + by_id, by_grp = {}, {} + for i in (0, 7, N - 1): + sql = f"SELECT id, name FROM {name} WHERE id = {i}" # nosec B608 - test-owned names + by_id[i] = [ + (r.get("id"), r.get("name")) for r in db.query("sql", sql).to_list() + ] + for g in range(5): + sql = f"SELECT id FROM {name} WHERE grp = {g}" # nosec B608 - test-owned names + by_grp[g] = sorted(r.get("id") for r in db.query("sql", sql).to_list()) + return by_id, by_grp + + +@pytest.fixture +def old_and_new(temp_db): + _type(temp_db, "Old") + _type(temp_db, "New") + _old_route(temp_db, "Old", _frame()) + return temp_db + + +def test_the_vertices_indexes_and_buckets_equal_those_of_graph_batch( + old_and_new, monkeypatch +): + db = old_and_new + with db.graph_batch() as batch: + graph_batch_class = type(batch) + + def refused(*args, **kwargs): + raise AssertionError( + "a vertex load through insert_dataframe went through GraphBatch.create_vertices" + ) + + monkeypatch.setattr(graph_batch_class, "create_vertices", refused) + assert db.insert_dataframe("New", _frame(), batch_size=10) == N + expected = _stored(db, "Old") + assert len(expected) == N and expected[7]["emb"][0] == "list" + assert expected[0].get("name", ("NoneType", None))[1] is None + assert sum(_bucket_counts(db, "Old")) == N and max(_bucket_counts(db, "Old")) < N + assert _stored(db, "New") == expected + assert _bucket_counts(db, "New") == _bucket_counts(db, "Old") + assert _index_view(db, "New") == _index_view(db, "Old") + assert db.lookup_by_key("New", ["id"], [7]).get("name") == "n7-é" + assert ( + db.query("sql", "SELECT outE().size() AS e FROM New WHERE id = 7") + .first() + .get("e") + == 0 + ) # a vertex + + +def test_insert_arrow_loads_a_vertex_type_the_same_way(old_and_new): + db = old_and_new + assert db.insert_arrow("New", pa.Table.from_pandas(_frame()), batch_size=16) == N + assert _stored(db, "New") == _stored(db, "Old") + + +def test_a_duplicate_key_fails_and_rolls_the_batch_back(temp_db): + _type(temp_db, "Dup") + frame = _frame(20) + frame.loc[15, "id"] = 3 # the unique index refuses it + with pytest.raises(ArcadeDBError, match="Duplicat"): + temp_db.insert_dataframe("Dup", frame, batch_size=50) + assert temp_db.is_transaction_active() is False + assert temp_db.query("sql", "SELECT count(*) AS n FROM Dup").first().get("n") == 0 + + +def test_a_vertex_load_refuses_to_run_inside_a_transaction_the_caller_opened(temp_db): + _type(temp_db, "Mine") + with temp_db.transaction(): + with pytest.raises( + ArcadeDBError, match="cannot run inside a transaction the caller opened" + ): + temp_db.insert_dataframe("Mine", _frame(5)) + assert temp_db.query("sql", "SELECT count(*) AS n FROM Mine").first().get("n") == 0 + assert temp_db.insert_dataframe("Mine", _frame(5)) == 5 # outside one it loads + + +def test_a_document_load_still_joins_the_callers_transaction(temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE Doc") + with pytest.raises(RuntimeError, match="stop"): + with temp_db.transaction(): + temp_db.insert_dataframe( + "Doc", pd.DataFrame({"id": [1, 2, 3]}), batch_size=2 + ) + raise RuntimeError("stop") + assert temp_db.query("sql", "SELECT count(*) AS n FROM Doc").first().get("n") == 0 diff --git a/tests/test_vertex_load_review_fixes.py b/tests/test_vertex_load_review_fixes.py new file mode 100644 index 00000000..a0eea0c4 --- /dev/null +++ b/tests/test_vertex_load_review_fixes.py @@ -0,0 +1,92 @@ +"""Review of the vertex loads through the columnar insert (#335): a commit conflict is retried, a failed later batch is exact. + +The old route, ``GraphBatch.create_vertices``, rolled a vertex chunk back and wrote it again (with a back-off, up to ``commitRetries``) +when its commit raised a ``NeedRetryException``: an MVCC conflict with another writer, a lock timeout. The columnar insert is one plain +transaction per chunk, so the same conflict would have surfaced as an error with earlier chunks already committed. The chunk loop of +``BulkOps.insertColumns`` (native wheel) and ``DocumentBatcher.insertColumns`` (JVM wheel) now does the retry, only in a transaction it +opened itself. Runs on both wheels. + +* several threads load one vertex type with a NOTUNIQUE index on a handful of key values, in many small batches: every row is stored; +* a duplicate key in a later batch: the earlier batches stay committed, the failing batch leaves no row and no index entry, the message + states the committed count, and the fixed tail loads afterwards. +""" + +import threading + +import numpy as np +import pyarrow as pa +import pytest + +from arcadedb.exceptions import ArcadeDBError + + +def _count(db, name): + sql = f"SELECT count(*) AS n FROM {name}" # nosec B608 - test-owned type name + return db.query("sql", sql).first().get("n") + + +def _vertex_type(db, name, unique=False): + db.command("sql", f"CREATE VERTEX TYPE {name} BUCKETS 3") + db.command("sql", f"CREATE PROPERTY {name}.id LONG") + db.command("sql", f"CREATE PROPERTY {name}.grp INTEGER") + db.command("sql", f"CREATE INDEX ON {name} (grp) NOTUNIQUE") + if unique: + db.command("sql", f"CREATE INDEX ON {name} (id) UNIQUE") + + +def _table(first, n, grps=5): + ids = np.arange(first, first + n, dtype=np.int64) + return pa.table({"id": ids, "grp": pa.array(ids % grps, pa.int32())}) + + +@pytest.mark.parametrize("threads", [2, 4]) +def test_threads_loading_one_vertex_type_in_small_batches_lose_no_row(temp_db, threads): + _vertex_type(temp_db, "V") + batches, batch_size = 20, 25 + errors = [] + + def load(worker): + try: + table = _table(worker * 1_000_000, batches * batch_size) + n = temp_db.insert_arrow("V", table, batch_size=batch_size) + assert n == batches * batch_size + except BaseException as e: # noqa: BLE001 - reported below, in the test thread + errors.append(repr(e)) + + pool = [threading.Thread(target=load, args=(w,)) for w in range(threads)] + for t in pool: + t.start() + for t in pool: + t.join(timeout=300) + assert not errors, errors + total = threads * batches * batch_size + assert _count(temp_db, "V") == total + distinct = temp_db.query("sql", "SELECT count(DISTINCT id) AS n FROM V") + assert distinct.first().get("n") == total # no row written twice by a retry + by_group = temp_db.query("sql", "SELECT grp, count(*) AS n FROM V GROUP BY grp") + assert sum(r.get("n") for r in by_group.to_list()) == total # the index agrees + + +def test_a_duplicate_in_a_later_batch_keeps_the_earlier_batches_and_nothing_of_its_own( + temp_db, +): + _vertex_type(temp_db, "Dup", unique=True) + ids = np.arange(500, dtype=np.int64) + ids[250] = 7 # the third batch of five repeats a key that the first batch committed + table = pa.table({"id": ids, "grp": pa.array(ids % 5, pa.int32())}) + with pytest.raises(ArcadeDBError, match="Duplicat") as err: + temp_db.insert_arrow("Dup", table, batch_size=100) + assert "the 200 row(s) of the batches before this one were committed" in str( + err.value + ) + assert temp_db.is_transaction_active() is False + assert _count(temp_db, "Dup") == 200 # batches 1 and 2, whole + for key in (200, 249, 251, 299): # the failing batch left no index entry behind + assert temp_db.lookup_by_key("Dup", ["id"], [key]) is None + assert temp_db.lookup_by_key("Dup", ["id"], [7]).get("id") == 7 # the committed one + + ids[250] = 250 # the fixed tail loads + tail = pa.table({"id": ids[200:], "grp": pa.array(ids[200:] % 5, pa.int32())}) + assert temp_db.insert_arrow("Dup", tail, batch_size=100) == 300 + assert _count(temp_db, "Dup") == 500 + assert temp_db.lookup_by_key("Dup", ["id"], [250]).get("id") == 250