From 68bc0a9a113ac5e3e71f556933185cf0a8d2bfe4 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Mon, 9 Nov 2026 18:58:29 +0900 Subject: [PATCH 1/6] Native: column frames for to_arrow, to_columns, to_dataframe and to_polars, read by the decoder the JVM wheel uses (refs #312) The native library answers a read of a cursor as one column frame of about 1 MiB (arcadedb_n_next_col, ColumnFrames.java), in the layout that ColumnBatcher writes for the JVM wheel: int64, float64, bool bytes, epoch milliseconds or microseconds, float32 vectors, Arrow strings (int32 offsets and a UTF-8 blob), DECIMAL as text, and a tagged JSON column for what no fixed layout carries without changing a value (nested values, records, a mixed column, an integer beyond 2**53 next to a float). arcadedb._columnar holds the one decoder and the merge of frames that both wheels call; results.py keeps only its JPype transport and the text of its json columns. The first reply of a query is still JSON for the row APIs, and the cursor keeps those rows (at most 256, until its next read by row) so that a columnar read replays them into its first frame. A result that the first reply held whole has no cursor: its dicts go through the old rules once, as one more chunk. The old JSON read path stays for to_list, to_json_list and row iteration in this commit. Tests: tests/test_column_frames.py (the decoder on hand-built frames, and the stored table through every columnar API, both wheels), tests/test_native_column_frames_reference.py (the old route copied into the test as the reference, native only). Co-Authored-By: Claude Sonnet 5.5 --- .../org/humemai/nativelib/ArcadeNative.java | 86 +++- .../org/humemai/nativelib/ColumnFrames.java | 382 +++++++++++++++ .../org/humemai/nativelib/NativeCodec.java | 7 + src/arcadedb/_columnar.py | 397 ++++++++++++++++ src/arcadedb/_native/colframes.py | 130 ++++++ src/arcadedb/_native/database.py | 96 +++- src/arcadedb/_native/lib.py | 27 ++ src/arcadedb/results.py | 408 ++-------------- tests/conftest.py | 2 + tests/test_column_frames.py | 338 ++++++++++++++ tests/test_native_column_frames_reference.py | 434 ++++++++++++++++++ 11 files changed, 1908 insertions(+), 399 deletions(-) create mode 100644 native/src/main/java/org/humemai/nativelib/ColumnFrames.java create mode 100644 src/arcadedb/_columnar.py create mode 100644 src/arcadedb/_native/colframes.py create mode 100644 tests/test_column_frames.py create mode 100644 tests/test_native_column_frames_reference.py diff --git a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java index 9b5acb91..52f9d3c5 100644 --- a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java +++ b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java @@ -20,6 +20,7 @@ import com.arcadedb.engine.ComponentFile; import com.arcadedb.exception.DuplicatedKeyException; import com.arcadedb.exception.NeedRetryException; +import com.arcadedb.query.sql.executor.Result; import com.arcadedb.query.sql.executor.ResultSet; import com.arcadedb.schema.VertexType; import com.arcadedb.server.ArcadeDBServer; @@ -36,6 +37,7 @@ import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -45,6 +47,17 @@ public final class ArcadeNative { static final Map DBS = new ConcurrentHashMap<>(); private static final Map RESULT_SETS = new ConcurrentHashMap<>(); + /** + * The rows of the first reply of an open cursor, kept until the cursor is read again or closed: that reply is JSON for the row APIs, + * and a columnar read (arcadedb_n_next_col) that starts from it replays these rows into its first frame instead of losing them. + * At most the first reply (256 rows) of a cursor that has more rows. + */ + private static final Map> HEADS = new ConcurrentHashMap<>(); + + private static ResultSet removeCursor(final long id) { + HEADS.remove(id); + return RESULT_SETS.remove(id); + } private static final Map SERVERS = new ConcurrentHashMap<>(); static final AtomicLong IDS = new AtomicLong(1); @@ -288,9 +301,10 @@ public static CCharPointer next(final IsolateThread thread, final long rsId, fin rs = RESULT_SETS.get(rsId); if (rs == null) throw new IllegalArgumentException("unknown result set " + rsId); - return out(batchOf(rsId, rs, max)); + HEADS.remove(rsId); + return out(batchOf(rsId, rs, max, null)); } catch (final Throwable t) { - RESULT_SETS.remove(rsId); + removeCursor(rsId); closeQuietly(rs); return err(t); } @@ -309,17 +323,65 @@ public static CCharPointer nextBin(final IsolateThread thread, final long rsId, if (rs == null) throw new IllegalArgumentException("unknown result set " + rsId); final ResultSet cursor = rs; - return withFloatTail(() -> batchOf(rsId, cursor, max)); + HEADS.remove(rsId); + return withFloatTail(() -> batchOf(rsId, cursor, max, null)); } catch (final Throwable t) { - RESULT_SETS.remove(rsId); + removeCursor(rsId); closeQuietly(rs); return err(t); } } + /** + * The next rows of an open cursor as one column frame (see ColumnFrames), written as an int64 LE byte count and the frame; a failure is + * the negated byte count of the error JSON and that JSON. max: rows at most (a frame also ends at about 1 MiB of values). columns: a + * JSON array of the column names to read, or "" for the union of the names of the rows. flags bit 0: datetimes as microseconds (the + * to_dataframe form). skip: of the rows that the first reply of this cursor already held (HEADS), the number the caller has taken; the + * first frame starts after them. The frame that ends the rows closes the cursor. The reply is freed with arcadedb_free. + */ + @CEntryPoint(name = "arcadedb_n_next_col") + @SuppressWarnings("unchecked") + public static CCharPointer nextCol(final IsolateThread thread, final long rsId, final int max, final CCharPointer columns, + final int flags, final int skip) { + ResultSet rs = null; + try { + rs = RESULT_SETS.get(rsId); + if (rs == null) + throw new IllegalArgumentException("unknown result set " + rsId); + final List head = HEADS.remove(rsId); + String[] pinned = null; + final String spec = str(columns); + if (spec != null && !spec.isEmpty()) { + final List names = (List) Json.parse(spec); + pinned = new String[names.size()]; + for (int i = 0; i < pinned.length; i++) + pinned[i] = (String) names.get(i); + } + final ColumnFrames.Frame frame = ColumnFrames.next(rs, head, skip, max, pinned, (flags & 1) != 0); + if (frame.done) { + removeCursor(rsId); + closeQuietly(rs); + } + return frameOut(frame.bytes.length, frame.bytes); + } catch (final Throwable t) { + removeCursor(rsId); + closeQuietly(rs); + final byte[] e = errorJson(t).getBytes(StandardCharsets.UTF_8); + return frameOut(-e.length, e); + } + } + + private static CCharPointer frameOut(final int length, final byte[] b) { + final CCharPointer p = UnmanagedMemory.malloc(8 + b.length); + final ByteBuffer bb = CTypeConversion.asByteBuffer(p, 8 + b.length).order(ByteOrder.LITTLE_ENDIAN); + bb.putLong(length); + bb.put(b, 0, b.length); + return p; + } + @CEntryPoint(name = "arcadedb_n_rs_close") public static CCharPointer rsClose(final IsolateThread thread, final long rsId) { - closeQuietly(RESULT_SETS.remove(rsId)); + closeQuietly(removeCursor(rsId)); return OK(); } @@ -328,16 +390,20 @@ private static String batch(final ResultSet rs, final int max) { return "{\"ok\":true,\"rs\":0,\"rows\":[],\"meta\":[],\"drained\":true}"; final long id = IDS.getAndIncrement(); RESULT_SETS.put(id, rs); - return batchOf(id, rs, max); + final List keep = new ArrayList<>(); + final String reply = batchOf(id, rs, max, keep); + if (RESULT_SETS.containsKey(id)) + HEADS.put(id, keep); + return reply; } - private static String batchOf(final long id, final ResultSet rs, final int max) { + private static String batchOf(final long id, final ResultSet rs, final int max, final List keep) { final StringBuilder sb = new StringBuilder("{\"ok\":true,\"rs\":"); final int mark = sb.length(); sb.append(id).append(','); - final boolean drained = NativeCodec.appendBatch(sb, rs, max); + final boolean drained = NativeCodec.appendBatch(sb, rs, max, keep); if (drained) { - RESULT_SETS.remove(id); + removeCursor(id); closeQuietly(rs); sb.replace(mark, mark + String.valueOf(id).length(), "0"); } @@ -650,7 +716,7 @@ public static CCharPointer lastError(final IsolateThread thread) { public static void rsCloseMany(final IsolateThread thread, final CLongPointer ids, final int n) { try { for (int i = 0; i < n; i++) - closeQuietly(RESULT_SETS.remove(ids.read(i))); + closeQuietly(removeCursor(ids.read(i))); } catch (final Throwable ignored) { // closing is best effort, as in arcadedb_n_rs_close } diff --git a/native/src/main/java/org/humemai/nativelib/ColumnFrames.java b/native/src/main/java/org/humemai/nativelib/ColumnFrames.java new file mode 100644 index 00000000..0c507346 --- /dev/null +++ b/native/src/main/java/org/humemai/nativelib/ColumnFrames.java @@ -0,0 +1,382 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Result rows as one column frame for the Python native backend (arcadedb_n_next_col): the layout the JVM wheel's ColumnBatcher + * writes, so one Python decoder (arcadedb/_columnar.py) reads both. A frame is + * + * int32 LE header length, header JSON, then for each column its null bitmap and its buffer + * + * The header is {"count":n,"done":bool,"cols":[{"name":..,"type":..,"nulls":bytes,"bytes":bytes[,"dim":d]},..]} ("done" is an addition + * to ColumnBatcher's header, which the decoder ignores). A null bitmap has bit i of byte i/8 set for a null row. A column type is + * + * i8 int64 LE (Long, Integer, Short, Byte) f8 float64 LE (Double, Float; a Float is its exact double) + * b1 one byte per row dt int64 epoch milliseconds, dtu: epoch microseconds (see micros) + * f4v dim float32 LE per row (float[]) f8v dim float64 LE per row (double[]) + * str int32 LE offsets (n + 1) then UTF-8 bytes dec a BigDecimal as text in the str layout + * json one JSON array of the column in the tagged encoding of NativeCodec: what no fixed layout carries without changing a value + * (nested lists and maps, records and RIDs, BigInteger, a column that mixes kinds, a DATE or a time-zone-aware temporal next to + * datetimes with micros, a ragged vector, a Long beyond 2**53 next to a floating value) + * + * A frame holds at most {@code max} rows and stops once the rows read add up to about FRAME_BYTES, so the working set of a read is a frame + * and not the table. + */ +package org.humemai.nativelib; + +import com.arcadedb.query.sql.executor.Result; +import com.arcadedb.query.sql.executor.ResultSet; + +import java.io.ByteArrayOutputStream; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.ArrayList; +import java.util.Calendar; +import java.util.Collection; +import java.util.Date; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +final class ColumnFrames { + /** The rows of a frame add up to about this many bytes of values. */ + static final long FRAME_BYTES = 1L << 20; + private static final long TWO53 = 1L << 53; + + private ColumnFrames() { + } + + /** One frame and whether the rows are over (the cursor was read to its end). */ + static final class Frame { + final byte[] bytes; + final boolean done; + + Frame(final byte[] bytes, final boolean done) { + this.bytes = bytes; + this.done = done; + } + } + + private static long toEpochMillis(final Object v) { + if (v instanceof Date d) + return d.getTime(); + if (v instanceof Calendar c) + return c.getTimeInMillis(); + if (v instanceof LocalDateTime ldt) + return ldt.toInstant(ZoneOffset.UTC).toEpochMilli(); + if (v instanceof LocalDate ld) + return ld.atStartOfDay().toInstant(ZoneOffset.UTC).toEpochMilli(); + if (v instanceof Instant i) + return i.toEpochMilli(); + if (v instanceof ZonedDateTime z) + return z.toInstant().toEpochMilli(); + if (v instanceof OffsetDateTime o) + return o.toInstant().toEpochMilli(); + throw new IllegalArgumentException("not a temporal: " + v.getClass()); + } + + private static long toEpochMicros(final Object v) { + if (v instanceof LocalDateTime ldt) + return Math.addExact(Math.multiplyExact(ldt.toEpochSecond(ZoneOffset.UTC), 1_000_000L), ldt.getNano() / 1000); + if (v instanceof Date d) + return Math.multiplyExact(d.getTime(), 1000L); + if (v instanceof Calendar c) + return Math.multiplyExact(c.getTimeInMillis(), 1000L); + throw new IllegalArgumentException("not a naive datetime: " + v.getClass()); + } + + private static boolean yearFits(final LocalDateTime ldt) { + return ldt.getYear() >= 1 && ldt.getYear() <= 9999; + } + + /** About how many bytes the values of a row take in a frame (a stop rule, not an exact size). */ + private static long estimate(final Result row) { + long n = 0; + for (final String name : row.getPropertyNames()) { + final Object v = row.getProperty(name); + if (v == null) + n += 1; + else if (v instanceof String s) + n += s.length() + 4; + else if (v instanceof float[] a) + n += 4L * a.length; + else if (v instanceof double[] a) + n += 8L * a.length; + else if (v instanceof Number || v instanceof Boolean || v instanceof Date || v instanceof LocalDate || v instanceof LocalDateTime + || v instanceof Instant) + n += 8; + else if (v instanceof Collection c) + n += 16 + 16L * c.size(); + else if (v instanceof Map m) + n += 16 + 32L * m.size(); + else + n += 32; + } + return n; + } + + /** + * The next frame. {@code head} holds rows that were already read from the cursor when its first reply was encoded as JSON (the + * frame starts at {@code head[skip]}, then goes on with the cursor), or is null. {@code pinned} is the column set, or null to take + * the union of the names of the rows of this frame, in order of first appearance. With {@code micros} a naive datetime column is + * epoch microseconds ("dtu") and a column with a DATE, a time-zone-aware temporal or a year Python cannot hold is "json". + */ + static Frame next(final ResultSet rs, final List head, final int skip, final int max, final String[] pinned, + final boolean micros) { + final List results = new ArrayList<>(Math.min(Math.max(max, 1), 1 << 14)); + int h = Math.max(skip, 0); + long budget = 0; + boolean done = false; + while (results.size() < max && budget < FRAME_BYTES) { + final Result row; + if (head != null && h < head.size()) + row = head.get(h++); + else if (rs.hasNext()) + row = rs.next(); + else { + done = true; + break; + } + results.add(row); + budget += estimate(row); + } + if (!done) + done = (head == null || h >= head.size()) && !rs.hasNext(); + + String[] columns = pinned; + if (columns == null) { + final Set names = new LinkedHashSet<>(); + String[] previous = null; + for (final Result row : results) { + final Set rowNames = row.getPropertyNames(); + if (previous != null && rowNames.size() == previous.length) { + int i = 0; + for (final String name : rowNames) + if (!name.equals(previous[i++])) { + i = -1; + break; + } + if (i >= 0) + continue; + } + names.addAll(rowNames); + previous = rowNames.toArray(new String[0]); + } + columns = names.toArray(new String[0]); + } + final int n = columns.length; + final int count = results.size(); + final List rows = new ArrayList<>(count); + for (final Result row : results) { + final Object[] vals = new Object[n]; + for (int c = 0; c < n; c++) + vals[c] = row.getProperty(columns[c]); + rows.add(vals); + } + + final StringBuilder hdr = new StringBuilder("{\"count\":").append(count).append(",\"done\":").append(done).append(",\"cols\":["); + final ByteArrayOutputStream body = new ByteArrayOutputStream(Math.max(count * n * 8, 16)); + + for (int c = 0; c < n; c++) { + final byte[] nulls = new byte[(count + 7) / 8]; + String type = null; + boolean sawInt = false, sawFloating = false, bigInt = false; + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (v == null) + continue; + final String t; + if (v instanceof Long || v instanceof Integer || v instanceof Short || v instanceof Byte) { + t = "i8"; + sawInt = true; + final long l = ((Number) v).longValue(); + if (l > TWO53 || l < -TWO53) + bigInt = true; + } else if (v instanceof Double || v instanceof Float) { + t = "f8"; + sawFloating = true; + } else if (v instanceof Boolean) + t = "b1"; + else if (v instanceof Date || v instanceof Calendar || v instanceof LocalDateTime || v instanceof LocalDate + || v instanceof Instant || v instanceof ZonedDateTime || v instanceof OffsetDateTime) { + if (!micros) + t = "dt"; + else if (v instanceof LocalDateTime ldt ? yearFits(ldt) : v instanceof Date || v instanceof Calendar) + t = "dtu"; + else + t = "json"; + } else if (v instanceof String || v instanceof Character) + t = "str"; + else if (v instanceof BigDecimal) + t = "dec"; + else if (v instanceof float[]) + t = "f4v"; + else if (v instanceof double[]) + t = "f8v"; + else + t = "json"; + if (type == null) + type = t; + else if (!type.equals(t)) { + if ((type.equals("i8") && t.equals("f8")) || (type.equals("f8") && t.equals("i8"))) + type = "f8"; + else { + type = "json"; + break; + } + } + } + if (type == null) + type = "str"; // an all-null column + // an integer beyond 2**53 next to a floating value has no exact double: the values stay what they are, in a json column + if (type.equals("f8") && sawInt && sawFloating && bigInt) + type = "json"; + + int vdim = -1; + if (type.equals("f4v") || type.equals("f8v")) { + for (int r = 0; r < count && vdim != -2; r++) { + final Object v = rows.get(r)[c]; + if (v == null) + continue; + final int len = type.equals("f4v") ? ((float[]) v).length : ((double[]) v).length; + if (vdim == -1) + vdim = len; + else if (vdim != len) + vdim = -2; + } + if (vdim < 0 && vdim != -1) + type = "json"; // ragged + else if (vdim == -1) + vdim = 0; + } + + final byte[] colBuf; + if (type.equals("i8") || type.equals("dt") || type.equals("dtu")) { + final ByteBuffer bb = ByteBuffer.allocate(count * 8).order(ByteOrder.LITTLE_ENDIAN); + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (v == null) { + nulls[r >> 3] |= (1 << (r & 7)); + bb.putLong(0); + } else if (type.equals("dt")) + bb.putLong(toEpochMillis(v)); + else if (type.equals("dtu")) + bb.putLong(toEpochMicros(v)); + else + bb.putLong(((Number) v).longValue()); + } + colBuf = bb.array(); + } else if (type.equals("f8")) { + final ByteBuffer bb = ByteBuffer.allocate(count * 8).order(ByteOrder.LITTLE_ENDIAN); + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (v == null) { + nulls[r >> 3] |= (1 << (r & 7)); + bb.putDouble(Double.NaN); + } else + bb.putDouble(((Number) v).doubleValue()); + } + colBuf = bb.array(); + } else if (type.equals("b1")) { + final ByteBuffer bb = ByteBuffer.allocate(count); + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (v == null) { + nulls[r >> 3] |= (1 << (r & 7)); + bb.put((byte) 0); + } else + bb.put((byte) (((Boolean) v) ? 1 : 0)); + } + colBuf = bb.array(); + } else if (type.equals("f4v")) { + final ByteBuffer bb = ByteBuffer.allocate(count * vdim * 4).order(ByteOrder.LITTLE_ENDIAN); + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (v == null) { + nulls[r >> 3] |= (1 << (r & 7)); + for (int k = 0; k < vdim; k++) + bb.putFloat(Float.NaN); + } else + for (final float x : (float[]) v) + bb.putFloat(x); + } + colBuf = bb.array(); + } else if (type.equals("f8v")) { + final ByteBuffer bb = ByteBuffer.allocate(count * vdim * 8).order(ByteOrder.LITTLE_ENDIAN); + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (v == null) { + nulls[r >> 3] |= (1 << (r & 7)); + for (int k = 0; k < vdim; k++) + bb.putDouble(Double.NaN); + } else + for (final double x : (double[]) v) + bb.putDouble(x); + } + colBuf = bb.array(); + } else if (type.equals("json")) { + final StringBuilder arr = new StringBuilder(); + arr.append('['); + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (r > 0) + arr.append(','); + if (v == null) + nulls[r >> 3] |= (1 << (r & 7)); + NativeCodec.append(arr, v); + } + arr.append(']'); + colBuf = arr.toString().getBytes(StandardCharsets.UTF_8); + } else { // str and dec: int32 offsets (count + 1) then the UTF-8 bytes + final ByteArrayOutputStream chars = new ByteArrayOutputStream(Math.max(count * 16, 16)); + final ByteBuffer offs = ByteBuffer.allocate((count + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + int pos = 0; + offs.putInt(0); + for (int r = 0; r < count; r++) { + final Object v = rows.get(r)[c]; + if (v == null) + nulls[r >> 3] |= (1 << (r & 7)); + else { + final String text = v instanceof BigDecimal d ? d.toPlainString() : v.toString(); + final byte[] b = text.getBytes(StandardCharsets.UTF_8); + chars.write(b, 0, b.length); + pos += b.length; + } + offs.putInt(pos); + } + final byte[] charBytes = chars.toByteArray(); + final ByteBuffer bb = ByteBuffer.allocate(offs.capacity() + charBytes.length); + bb.put(offs.array()); + bb.put(charBytes); + colBuf = bb.array(); + } + + if (c > 0) + hdr.append(','); + hdr.append("{\"name\":"); + NativeCodec.quote(hdr, columns[c]); + hdr.append(",\"type\":\"").append(type).append("\",\"nulls\":").append(nulls.length).append(",\"bytes\":").append(colBuf.length); + if (vdim > 0 && (type.equals("f4v") || type.equals("f8v"))) + hdr.append(",\"dim\":").append(vdim); + hdr.append('}'); + body.write(nulls, 0, nulls.length); + body.write(colBuf, 0, colBuf.length); + } + hdr.append("]}"); + + final byte[] headerBytes = hdr.toString().getBytes(StandardCharsets.UTF_8); + final byte[] bodyBytes = body.toByteArray(); + final ByteBuffer out = ByteBuffer.allocate(4 + headerBytes.length + bodyBytes.length).order(ByteOrder.LITTLE_ENDIAN); + out.putInt(headerBytes.length); + out.put(headerBytes); + out.put(bodyBytes); + return new Frame(out.array(), done); + } +} diff --git a/native/src/main/java/org/humemai/nativelib/NativeCodec.java b/native/src/main/java/org/humemai/nativelib/NativeCodec.java index 59b7c94c..5d037a50 100644 --- a/native/src/main/java/org/humemai/nativelib/NativeCodec.java +++ b/native/src/main/java/org/humemai/nativelib/NativeCodec.java @@ -52,6 +52,11 @@ private NativeCodec() { * for a projection, or [rid, typeName, kind] for a record row (kind v, e or d). */ static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int max) { + return appendBatch(sb, rs, max, null); + } + + /** As above; {@code keep}, when not null, receives each row that was encoded (a columnar read replays them, see ArcadeNative.HEADS). */ + static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int max, final java.util.List keep) { final StringBuilder meta = new StringBuilder("["); sb.append("\"rows\":["); int n = 0; @@ -78,6 +83,8 @@ static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int throw t; } n++; + if (keep != null) + keep.add(row); if (sb.length() >= MAX_BATCH_CHARS) break; } diff --git a/src/arcadedb/_columnar.py b/src/arcadedb/_columnar.py new file mode 100644 index 00000000..7cb62e36 --- /dev/null +++ b/src/arcadedb/_columnar.py @@ -0,0 +1,397 @@ +"""The column frame of both wheels, and the one decoder that reads it. + +A frame is what ``ColumnBatcher`` (the JVM wheel's bridge jar) writes and what the native library's ``arcadedb_n_next_col`` writes: a +little-endian ``int32`` header length, a JSON header (``count``, ``cols``: name, type, null-bitmap bytes, data bytes and, for a +vector, ``dim``), then for each column its null bitmap (bit i of byte i/8 set means row i is null) and its buffer. A column type is +``i8`` (int64), ``f8`` (float64), ``b1`` (one byte per bool), ``dt`` and ``dtu`` (int64 epoch milliseconds and microseconds), +``f4v`` and ``f8v`` (``dim`` floats per row), ``str`` (int32 offsets, count + 1 of them, then the UTF-8 bytes: Arrow's own string +layout), ``dec`` (a BigDecimal as text in the string layout) and ``json`` (a JSON array of the column, for values no fixed layout +carries: nested lists and maps, RIDs, a DATE next to datetimes, a mixed column). + +The interface and the format are shared, the plumbing is per wheel: what differs is how the bytes get here (JPype or ctypes) and the +text of a ``json`` column (the JVM wheel's tagged JSON with a side array of Java objects, the native library's tagged JSON), so the +decoders take ``json_values(data, side, count)`` from the wheel. Nothing here imports JPype. +""" + +import json +from typing import Any, Callable, Dict, List, Tuple + + +def _cast_all_to_string(arrs: list, pa) -> list: + """ + Cast every array to a string type. + + Tolerates a type pyarrow cannot cast directly (e.g. a vector/list column) + by going through Python objects instead. + """ + unified = [] + for a in arrs: + try: + unified.append(a.cast(pa.string())) + except pa.ArrowException: + unified.append( + pa.array( + [None if v is None else str(v) for v in a.to_pylist()], + type=pa.string(), + ) + ) + return unified + + +def _decode_strings(data, count: int, mask) -> list: + """The strings of one column buffer: int32 offsets, then the UTF-8 blob. + + ``mask`` is the null mask (True = null) or None when the column has no nulls. + """ + import numpy as np + + offs = np.frombuffer(data[: (count + 1) * 4], dtype=" list: + """The DECIMAL values of one column buffer, as ``Decimal`` objects (None for + null). The bridge writes each ``BigDecimal`` in full as text, so no digit is + lost to a double.""" + from decimal import Decimal + + return [ + None if s is None else Decimal(s) for s in _decode_strings(data, count, mask) + ] + + +def _missing_column_part(count: int, template): + """A batch of ``count`` rows that lack a column some other batch has: nulls in + the shape ``to_columns`` uses for that column (NaN for numbers, NaT for + datetimes, None for everything else).""" + import numpy as np + + if isinstance(template, np.ndarray): + if template.ndim == 2: + return np.full((count, template.shape[1]), np.nan, dtype=template.dtype) + if template.dtype.kind in "iuf": + return np.full(count, np.nan) + if template.dtype.kind == "M": + return np.full(count, np.datetime64("NaT"), dtype=template.dtype) + if template.dtype.kind == "O": + return np.full(count, None, dtype=object) + return [None] * count + + +def _is_untyped_arrow_chunk(a, pa) -> bool: + """True for a chunk that says nothing about its column's type: all values + null (a batch whose rows lack the property, or carry it as null), or a list + column holding only empty lists, which Arrow types ``list``.""" + if len(a) == 0 or a.null_count == len(a) or pa.types.is_null(a.type): + return True + return pa.types.is_list(a.type) and pa.types.is_null(a.type.value_type) + + +def _common_decimal_type(types: list, pa): + """One decimal type every given decimal type fits in (decimal128 up to 38 + digits, decimal256 up to 76), or None when it would need more.""" + scale = max(t.scale for t in types) + precision = max(t.precision - t.scale for t in types) + scale + if precision <= 38: + return pa.decimal128(precision, scale) + if precision <= 76: + return pa.decimal256(precision, scale) + return None + + +def _unify_arrow_chunk_types(arrs: list, pa) -> list: + """ + Unify a column's per-batch Arrow arrays onto one Arrow type. + + ArcadeDB is schemaless per document, so a property's Java type can + legally vary row to row; ColumnBatcher infers each batch's column type + independently, so two batches of the same result set can produce + incompatible Arrow types for the same column (#7108) - e.g. int64 in one + batch, float64 in the next. ``pa.chunked_array`` requires every chunk to + share one type, so leaving them as-is raises at concatenation time on a + result set that is otherwise entirely legal. + + Numeric types widen (int64 -> float64, matching to_columns' pandas-style + promotion) when every value survives the cast exactly; pyarrow's cast is + safe by default and raises rather than silently lose precision, which an + int64 outside float64's +-2**53 exact range would (code review on #7108), + so that case - and anything else, including a json or vector column that + varied - degrades to string instead, so the column always ends up with + exactly one type rather than an exception. + + int64/float64 is the only numeric pair worth special-casing: ColumnBatcher + (src/java/.../ColumnBatcher.java) already widens every + Java integer type (Long/Integer/Short/Byte) to the same "i8" wire encoding + and every Java floating type (Double/Float) to "f8" *within* one batch, so + across batches the only numeric type mismatch this ever has to unify is + exactly this pair - there is no int32/int64 or bool/int64 case to widen. + """ + types = {a.type for a in arrs} + if len(types) <= 1: + return arrs + + # A chunk that carries no type information (every value null, or only empty + # lists) takes the type the other chunks agree on, so the column's type does + # not depend on where the batch boundaries fall (#114). + untyped = [_is_untyped_arrow_chunk(a, pa) for a in arrs] + if any(untyped) and not all(untyped): + typed = _unify_arrow_chunk_types( + [a for a, u in zip(arrs, untyped) if not u], pa + ) + target = typed[0].type + filled = iter(typed) + out = [] + for a, u in zip(arrs, untyped): + if not u: + out.append(next(filled)) + elif a.null_count == len(a): + out.append(pa.nulls(len(a), type=target)) + else: + try: + out.append(a.cast(target)) + except pa.ArrowException: + return _cast_all_to_string(arrs, pa) + return out + + if types <= {pa.int64(), pa.float64()}: + try: + return [a.cast(pa.float64()) for a in arrs] + except pa.ArrowException: + pass + + if all(pa.types.is_decimal(t) for t in types): + common = _common_decimal_type(list(types), pa) + if common is not None: + return [a.cast(common) for a in arrs] + + return _cast_all_to_string(arrs, pa) + + +def _null_mask(buf, pos: int, nulls_len: int, count: int): + """``(null_bits, mask)`` of a column: ``mask`` is a bool array (True = null) or None when no row is null.""" + import numpy as np + + null_bits = np.frombuffer(buf[pos : pos + nulls_len], dtype=np.uint8) + if not null_bits.any(): + return null_bits, None + return null_bits, np.unpackbits(null_bits, bitorder="little")[:count].astype(bool) + + +def read_header(buf) -> Tuple[Dict[str, Any], int]: + """``(header, position of the first column)`` of a frame.""" + hlen = int.from_bytes(buf[:4], "little") + return json.loads(bytes(buf[4 : 4 + hlen])), 4 + hlen + + +def decode_columns_frame( + buf, side, json_values: Callable, micros: bool = False +) -> Tuple[int, Dict[str, Any], Dict[str, Any]]: + """``(count, parts, header)`` of one frame as ``to_columns`` columns: ``parts`` maps a column name to a numpy array + (int64, float64, bool, datetime64) or a Python list (text, JSON-typed values, a bool column with nulls); a count of 0 has + no parts. With ``micros`` a naive datetime column is ``datetime64[us]`` (``dtu``). + """ + import numpy as np + + header, pos = read_header(buf) + count = header["count"] + if count == 0: + return 0, {}, header + parts: Dict[str, Any] = {} + for col in header["cols"]: + name, ctype = col["name"], col["type"] + nulls_len = col["nulls"] + _, mask = _null_mask(buf, pos, nulls_len, count) + has_nulls = mask is not None + pos += nulls_len + data = buf[pos : pos + col["bytes"]] + pos += col["bytes"] + + if ctype == "i8": + arr = np.frombuffer(data, dtype=" 2-D array (count, dim); the writer puts NaN rows into null slots; + # an empty vector has no "dim" in the header + dim = col.get("dim", 0) + dt = " Dict[str, Any]: + """The ``to_columns`` dict of the parts of every frame: a column a frame lacks gets nulls (NaN, NaT or None) there.""" + import numpy as np + + out: Dict[str, Any] = {} + for name in names: + template = next(p[name] for _, p in batches if name in p) + parts = [ + p[name] if name in p else _missing_column_part(c, template) + for c, p in batches + ] + np_parts = [p for p in parts if not isinstance(p, list)] + if parts and len(np_parts) == len(parts): + try: + out[name] = np.concatenate(parts) if len(parts) > 1 else parts[0] + continue + except Exception: # nosec B110 - fall through to generic merge + pass + column: list = [] + for p in parts: + column.extend(p if isinstance(p, list) else p.tolist()) + out[name] = column + return out + + +def decode_arrow_frame( + buf, side, json_values: Callable +) -> Tuple[int, Dict[str, Any], Dict[str, Any]]: + """``(count, chunks, header)`` of one frame as ``to_arrow`` columns: ``chunks`` maps a column name to a ``pyarrow.Array``. + Integers keep their type and their nulls (``to_columns`` would make them float64 with NaN), strings wrap the frame's own + offsets and bytes, vectors are fixed-size lists of float32 or float64.""" + import numpy as np + import pyarrow as pa + + header, pos = read_header(buf) + count = header["count"] + if count == 0: + return 0, {}, header + chunks: Dict[str, Any] = {} + for col in header["cols"]: + name, ctype = col["name"], col["type"] + nulls_len = col["nulls"] + null_bits, mask = _null_mask(buf, pos, nulls_len, count) + has_nulls = mask is not None + pos += nulls_len + data = buf[pos : pos + col["bytes"]] + pos += col["bytes"] + + if ctype == "i8": + # kept as int64 WITH validity, where to_columns promotes the whole column to float64/NaN + arr = pa.array(np.frombuffer(data, dtype="`` array of a column of float lists that float32 holds, or None.""" + import pyarrow as pa + + vec = _rows._float_vectors(values) + if vec is None: + return None + arr = pa.FixedSizeListArray.from_arrays(pa.array(vec.reshape(-1)), vec.shape[1]) + if any(v is None for v in values): + arr = pa.array(arr.to_pylist(), type=arr.type, mask=[v is None for v in values]) + return arr + + +def arrow_result(head_rows, columns, batches): + """The ``pyarrow.Table`` of the buffered rows (dicts) followed by the decoded frames ``batches`` (a list of ``(count, chunks)``).""" + import pyarrow as pa + + names = [] + chunks_list = [] + if head_rows: + table = _rows.to_arrow(head_rows, columns) + chunks = {} + for name in table.column_names: + column = table.column(name) + arr = ( + column.chunk(0) + if column.num_chunks == 1 + else pa.concat_arrays(column.chunks) + ) + if pa.types.is_null(arr.type): + # a frame types an all-null column as text + arr = pa.nulls(len(arr), pa.string()) + elif pa.types.is_list(arr.type) and arr.type.value_type == pa.float64(): + other = _first_part(batches, name) + if other is None or pa.types.is_fixed_size_list(other.type): + vector = _vector_arrow([row.get(name) for row in head_rows]) + if vector is not None: + arr = vector + chunks[name] = arr + names.extend(chunks) + chunks_list.append((len(head_rows), chunks)) + for count, chunks in batches: + for name in chunks: + if name not in names: + names.append(name) + chunks_list.append((count, chunks)) + return _columnar.merge_arrow(names, chunks_list) diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index f30dc526..71d178eb 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -210,6 +210,9 @@ def __init__(self, database, reply: dict, tagged: bool = True): self._rs = reply["rs"] self._buf = deque(zip(reply["rows"], reply["meta"])) self._more = not reply["drained"] + # the rows of this reply that the cursor still holds, for a columnar read to replay (arcadedb_n_next_col): the first reply of + # a cursor with more rows. Any later read of the cursor by row (arcadedb_n_next) lets the library drop them + self._head_n = len(reply["rows"]) if self._more else 0 self._closed = False self._exhausted = False # the engine computes lazily: an error met while the first rows were read is raised when the rows before it @@ -224,6 +227,7 @@ def _lazy(cls, database, rs_id: int) -> "NativeResultSet": self._rs = rs_id self._buf = deque() self._tagged = False + self._head_n = 0 self._more = rs_id > 0 self._closed = False self._exhausted = False @@ -258,6 +262,7 @@ def _fill(self) -> bool: raise _read_error(err) if not self._more: return False + self._head_n = 0 try: if lib.NEXT_BIN and _BINARY_REPLY: reply, tagged = loads_tagged( @@ -407,33 +412,100 @@ def iter_json_batches(self, batch_size: int = 10_000, **kwargs): # returns as it is. The flag is read after the batch is built, so it covers every reply the batch's rows came from yield [json_native(row) for row in batch] if self._tagged else batch + def _read_frames(self, batch_size, columns, decode, flags=0): + """The rows still to be read, in the two forms a result has: ``(head_rows, batches)``. + + ``head_rows`` are rows that this object holds as dicts (the first reply, when the cursor no longer has them); + ``batches`` is the list of ``(count, parts)`` that ``decode(frame_bytes)`` makes of each column frame the cursor + yields (about 1 MiB of values each, ``batch_size`` rows at most; ``flags`` 1 asks for datetimes as microseconds). The result + is finished when this returns. + """ + if int(batch_size) < 1: + raise ValueError("batch_size must be at least 1") + head_rows: list = [] + batches: list = [] + if not self._readable(): + return head_rows, batches + spec = json.dumps(list(columns)).encode("utf-8") if columns is not None else b"" + skip = 0 + if self._head_n and self._more and self._rs: + # the cursor still holds the rows of the first reply: it replays the ones this object has not handed out + skip = self._head_n - len(self._buf) + self._buf.clear() + else: + head_rows = [props for props, _ in self._buf] + self._buf.clear() + if self._pending_error is not None: + self._fill() # the rows before the error are consumed: raise it + while self._more and self._rs: + try: + raw = lib.call_frame(self._rs, int(batch_size), spec, flags, skip) + except lib.ColumnFrameError as exc: + # the library closed the cursor + self._more = False + self._rs = 0 + self._closed = True + self._exhausted = True + raise _read_error(exc.reply) from None + skip = 0 + self._head_n = 0 + count, parts, header = decode(raw) + if header.get("done"): + self._more = False # the library closed the cursor with the frame that ended the rows + self._rs = 0 + if count: + batches.append((count, parts)) + self._finish() + return head_rows, batches + def to_dataframe(self, convert_types: bool = True): + """The rows as a ``pandas.DataFrame``: the values ``to_list()`` returns (a DATETIME keeps its microseconds, a DATE is a + ``datetime.date`` in an object column, a ``Decimal`` stays one), read as column frames, not as one dict per row. + """ import pandas as pd - return pd.DataFrame(self._drain_dicts()) + columns = self._columns(25_000, None, True) + # pandas rejects 2-D ndarrays as column values; vector columns become object columns of row slices + columns = { + k: (list(v) if getattr(v, "ndim", 1) > 1 else v) for k, v in columns.items() + } + return pd.DataFrame(columns) + + def _columns(self, batch_size, columns, micros): + from . import colframes + + head_rows, batches = self._read_frames( + batch_size, + columns, + lambda raw: colframes.decode_columns(raw, micros), + 1 if micros else 0, + ) + return colframes.columns_result(head_rows, columns, micros, batches) def to_columns(self, batch_size: int = 25_000, columns=None): """Columns as the JPype path returns them (numpy int64/float64/bool/datetime64[ms] arrays, lists for text and - JSON-typed values, an object array of Decimal), built in Python from the rows: a columnar JSON path, no - binary buffer. None when numpy is missing.""" + JSON-typed values, an object array of Decimal), decoded from column frames by the decoder both builds share. None when + numpy is missing.""" try: - import numpy # noqa: F401 + __import__("numpy") # availability check except ImportError: return None - from .columns import to_columns - - return to_columns(self._drain_dicts(), columns) + return self._columns(batch_size, columns, False) def to_arrow(self, batch_size: int = 25_000, columns=None): - """A ``pyarrow.Table`` with the types of :meth:`to_columns` but nulls kept (None when pyarrow is missing).""" + """A ``pyarrow.Table`` with the types of :meth:`to_columns` but nulls kept (None when pyarrow is missing), decoded from + column frames by the decoder both builds share.""" try: - import numpy # noqa: F401 - import pyarrow # noqa: F401 + __import__("numpy") # availability check + __import__("pyarrow") except ImportError: return None - from .columns import to_arrow + from . import colframes - return to_arrow(self._drain_dicts(), columns) + head_rows, batches = self._read_frames( + batch_size, columns, colframes.decode_arrow + ) + return colframes.arrow_result(head_rows, columns, batches) def to_polars(self, batch_size: int = 25_000, columns=None): try: diff --git a/src/arcadedb/_native/lib.py b/src/arcadedb/_native/lib.py index 58b261ba..dc3f1e98 100644 --- a/src/arcadedb/_native/lib.py +++ b/src/arcadedb/_native/lib.py @@ -232,6 +232,7 @@ class _Params(ctypes.Structure): "arcadedb_n_insert_colp": [P, L, S, S, P, I], "arcadedb_n_next": [P, L, I], "arcadedb_n_next_bin": [P, L, I], + "arcadedb_n_next_col": [P, L, I, S, I, I], "arcadedb_n_rs_close": [P, L], "arcadedb_n_cmd": [P, L, I, S, S, S, S, L, I], "arcadedb_n_last_error": [P], @@ -365,6 +366,32 @@ def call_bin(name: str, *args): return raw.decode("utf-8"), floats +def call_frame(rs_id: int, max_rows: int, spec: bytes, flags: int, skip: int) -> bytes: + """One column frame of an open cursor (``arcadedb_n_next_col``): its bytes, in the layout ``arcadedb._columnar`` reads. + + The reply is an int64 byte count and the frame. A negative count is the negated length of an error JSON: it is raised as + :class:`ColumnFrameError` with the decoded reply (the cursor is already closed on the Java side then). + """ + lib = _lib or _load() + th = _thread_handle() + ptr = lib.arcadedb_n_next_col(th, rs_id, max_rows, spec, flags, skip) + try: + n = int.from_bytes(ctypes.string_at(ptr, 8), "little", signed=True) + if n < 0: + raise ColumnFrameError(json.loads(ctypes.string_at(ptr + 8, -n))) + return ctypes.string_at(ptr + 8, n) + finally: + lib.arcadedb_free(th, ptr) + + +class ColumnFrameError(Exception): + """The error reply (the dict ``errorJson`` writes) of a failed ``arcadedb_n_next_col``.""" + + def __init__(self, reply: dict): + super().__init__(reply.get("error")) + self.reply = reply + + def alloc(size: int): """Memory from the library's allocator (``arcadedb_n_alloc``), for the reply of a callback. diff --git a/src/arcadedb/results.py b/src/arcadedb/results.py index ebb344d0..945ab21e 100644 --- a/src/arcadedb/results.py +++ b/src/arcadedb/results.py @@ -12,6 +12,12 @@ import jpype as _jpype from jpype import JException +from ._columnar import ( + decode_arrow_frame, + decode_columns_frame, + merge_arrow, + merge_columns, +) from ._logging import get_logger from .callbacks import raise_pending_interrupt from .exceptions import ArcadeDBError @@ -92,6 +98,11 @@ def hook(obj): return hook +def _typed_json_values(data, side, count): + """The values of a ``json`` column of a ColumnBatcher frame: the tagged JSON (TypedRows) with its side array of Java objects.""" + return json.loads(bytes(data), object_hook=_typed_hook(side)) + + def _loads_typed_row(java_row): """The dict of one row or record that ``ReaderCalls.rowJson`` / ``documentJson`` wrote. @@ -156,162 +167,6 @@ def _java_class_name(value: Any) -> str: return str(value.getClass().getName()) -def _cast_all_to_string(arrs: list, pa) -> list: - """ - Cast every array to a string type. - - Tolerates a type pyarrow cannot cast directly (e.g. a vector/list column) - by going through Python objects instead. - """ - unified = [] - for a in arrs: - try: - unified.append(a.cast(pa.string())) - except pa.ArrowException: - unified.append( - pa.array( - [None if v is None else str(v) for v in a.to_pylist()], - type=pa.string(), - ) - ) - return unified - - -def _decode_strings(data, count: int, mask) -> list: - """The strings of one column buffer: int32 offsets, then the UTF-8 blob. - - ``mask`` is the null mask (True = null) or None when the column has no nulls. - """ - import numpy as np - - offs = np.frombuffer(data[: (count + 1) * 4], dtype=" list: - """The DECIMAL values of one column buffer, as ``Decimal`` objects (None for - null). The bridge writes each ``BigDecimal`` in full as text, so no digit is - lost to a double.""" - from decimal import Decimal - - return [ - None if s is None else Decimal(s) for s in _decode_strings(data, count, mask) - ] - - -def _missing_column_part(count: int, template): - """A batch of ``count`` rows that lack a column some other batch has: nulls in - the shape ``to_columns`` uses for that column (NaN for numbers, NaT for - datetimes, None for everything else).""" - import numpy as np - - if isinstance(template, np.ndarray): - if template.ndim == 2: - return np.full((count, template.shape[1]), np.nan, dtype=template.dtype) - if template.dtype.kind in "iuf": - return np.full(count, np.nan) - if template.dtype.kind == "M": - return np.full(count, np.datetime64("NaT"), dtype=template.dtype) - if template.dtype.kind == "O": - return np.full(count, None, dtype=object) - return [None] * count - - -def _is_untyped_arrow_chunk(a, pa) -> bool: - """True for a chunk that says nothing about its column's type: all values - null (a batch whose rows lack the property, or carry it as null), or a list - column holding only empty lists, which Arrow types ``list``.""" - if len(a) == 0 or a.null_count == len(a) or pa.types.is_null(a.type): - return True - return pa.types.is_list(a.type) and pa.types.is_null(a.type.value_type) - - -def _common_decimal_type(types: list, pa): - """One decimal type every given decimal type fits in (decimal128 up to 38 - digits, decimal256 up to 76), or None when it would need more.""" - scale = max(t.scale for t in types) - precision = max(t.precision - t.scale for t in types) + scale - if precision <= 38: - return pa.decimal128(precision, scale) - if precision <= 76: - return pa.decimal256(precision, scale) - return None - - -def _unify_arrow_chunk_types(arrs: list, pa) -> list: - """ - Unify a column's per-batch Arrow arrays onto one Arrow type. - - ArcadeDB is schemaless per document, so a property's Java type can - legally vary row to row; ColumnBatcher infers each batch's column type - independently, so two batches of the same result set can produce - incompatible Arrow types for the same column (#7108) - e.g. int64 in one - batch, float64 in the next. ``pa.chunked_array`` requires every chunk to - share one type, so leaving them as-is raises at concatenation time on a - result set that is otherwise entirely legal. - - Numeric types widen (int64 -> float64, matching to_columns' pandas-style - promotion) when every value survives the cast exactly; pyarrow's cast is - safe by default and raises rather than silently lose precision, which an - int64 outside float64's +-2**53 exact range would (code review on #7108), - so that case - and anything else, including a json or vector column that - varied - degrades to string instead, so the column always ends up with - exactly one type rather than an exception. - - int64/float64 is the only numeric pair worth special-casing: ColumnBatcher - (src/java/.../ColumnBatcher.java) already widens every - Java integer type (Long/Integer/Short/Byte) to the same "i8" wire encoding - and every Java floating type (Double/Float) to "f8" *within* one batch, so - across batches the only numeric type mismatch this ever has to unify is - exactly this pair - there is no int32/int64 or bool/int64 case to widen. - """ - types = {a.type for a in arrs} - if len(types) <= 1: - return arrs - - # A chunk that carries no type information (every value null, or only empty - # lists) takes the type the other chunks agree on, so the column's type does - # not depend on where the batch boundaries fall (#114). - untyped = [_is_untyped_arrow_chunk(a, pa) for a in arrs] - if any(untyped) and not all(untyped): - typed = _unify_arrow_chunk_types( - [a for a, u in zip(arrs, untyped) if not u], pa - ) - target = typed[0].type - filled = iter(typed) - out = [] - for a, u in zip(arrs, untyped): - if not u: - out.append(next(filled)) - elif a.null_count == len(a): - out.append(pa.nulls(len(a), type=target)) - else: - try: - out.append(a.cast(target)) - except pa.ArrowException: - return _cast_all_to_string(arrs, pa) - return out - - if types <= {pa.int64(), pa.float64()}: - try: - return [a.cast(pa.float64()) for a in arrs] - except pa.ArrowException: - pass - - if all(pa.types.is_decimal(t) for t in types): - common = _common_decimal_type(list(types), pa) - if common is not None: - return [a.cast(common) for a in arrs] - - return _cast_all_to_string(arrs, pa) - - class ResultSet: """Iterator wrapper for ArcadeDB query results. @@ -816,98 +671,20 @@ def _read_columns(self, batch_size, columns, micros): column of ``date`` objects, which is what ``to_dataframe()`` needs to return the values ``to_list()`` does.""" try: - import numpy as np + __import__("numpy") # availability check except ImportError: return None - import json - column_batcher = _bridge_class("ColumnBatcher") if column_batcher is None: return None # Every batch reports the union of its own rows' property names. The - # result's columns are the union over the batches, in order of first - # appearance (#113); a batch lacking one gets nulls for it when merged. + # result's columns are the union over the batches, in order of + # first appearance (#113); a batch lacking one gets nulls for it when merged. names: List[str] = [] batches: List[Tuple[int, Dict[str, Any]]] = [] - def decode_batch(buf, side): - hlen = int.from_bytes(buf[:4], "little") - header = json.loads(bytes(buf[4 : 4 + hlen])) - count = header["count"] - if count == 0: - return 0 - batch_parts: Dict[str, Any] = {} - pos = 4 + hlen - for col in header["cols"]: - name, ctype = col["name"], col["type"] - nulls_len = col["nulls"] - null_bits = np.frombuffer(buf[pos : pos + nulls_len], dtype=np.uint8) - has_nulls = bool(null_bits.any()) - if has_nulls: - mask = np.unpackbits(null_bits, bitorder="little")[:count].astype( - bool - ) - pos += nulls_len - data = buf[pos : pos + col["bytes"]] - pos += col["bytes"] - - if ctype == "i8": - arr = np.frombuffer(data, dtype=" 2-D array (count, dim); - # Java already writes NaN rows into null slots - # an empty vector has no "dim" in the header - dim = col.get("dim", 0) - dt = " 1 else parts[0] - continue - except Exception: # nosec B110 - fall through to generic merge - pass - column = [] - for p in parts: - column.extend(p if isinstance(p, list) else p.tolist()) - out[name] = column - return out + return merge_columns(names, batches) def to_arrow( self, batch_size: int = 25_000, columns: Optional[Sequence[str]] = None @@ -983,13 +747,11 @@ def to_arrow( falls back to the row-based paths. """ try: - import numpy as np + __import__("numpy") # availability check import pyarrow as pa except ImportError: return None - import json - column_batcher = _bridge_class("ColumnBatcher") if column_batcher is None: return None @@ -1000,107 +762,6 @@ def to_arrow( names: List[str] = [] batches: List[Tuple[int, Dict[str, Any]]] = [] - def validity(null_bits, count, has_nulls): - """Arrow validity is 1=valid; the bridge writes 1=null, so invert. - Bits past `count` are padding and Arrow ignores them.""" - if not has_nulls: - return None - return pa.py_buffer(bytes(np.invert(null_bits).tobytes())) - - def decode_batch(buf, side): - hlen = int.from_bytes(buf[:4], "little") - header = json.loads(bytes(buf[4 : 4 + hlen])) - count = header["count"] - if count == 0: - return 0 - batch_chunks: Dict[str, Any] = {} - pos = 4 + hlen - for col in header["cols"]: - name, ctype = col["name"], col["type"] - nulls_len = col["nulls"] - null_bits = np.frombuffer(buf[pos : pos + nulls_len], dtype=np.uint8) - has_nulls = bool(null_bits.any()) - mask = None - if has_nulls: - mask = np.unpackbits(null_bits, bitorder="little")[:count].astype( - bool - ) - pos += nulls_len - data = buf[pos : pos + col["bytes"]] - pos += col["bytes"] - - if ctype == "i8": - # kept as int64 WITH validity, where to_columns would have - # promoted the whole column to float64/NaN - arr = pa.array(np.frombuffer(data, dtype=" Python upcalls (Python functions in SQL) exist on the native build only collect_ignore.append("test_native_udf.py") # how find_nearest hands the query vector to the native library (the codec's side buffer) diff --git a/tests/test_column_frames.py b/tests/test_column_frames.py new file mode 100644 index 00000000..a2e51a22 --- /dev/null +++ b/tests/test_column_frames.py @@ -0,0 +1,338 @@ +"""The column frame decoder both wheels share (``arcadedb._columnar``), and the columnar readers that use it (refs #312). + +The first half builds frames by hand, in the layout of ``ColumnBatcher`` (JVM wheel) and ``arcadedb_n_next_col`` (native wheel), and +decodes them: no database, so a decoder change is checked on its own. The second half reads one stored table through ``to_arrow``, +``to_columns``, ``to_dataframe`` and ``to_polars`` and compares each value with what was stored, on whichever wheel runs the tests: +nulls kept in Arrow, one type per column whatever the batch size, a result of more than one frame, column subsets, an empty result, +one row, many columns, vectors of several lengths. There is no branch on ``arcadedb.__backend__``: a literal that holds on both wheels +is the parity check. +""" + +import json +import math +import struct +from datetime import datetime + +import numpy as np +import pyarrow as pa +import pytest + +from arcadedb import _columnar + + +def _json_values(data, side, count): + return json.loads(bytes(data)) + + +def _bits(count, nulls): + out = bytearray((count + 7) // 8) + for i in nulls: + out[i >> 3] |= 1 << (i & 7) + return bytes(out) + + +def _frame(count, columns): + """``columns``: ``(name, type, null_rows, data_bytes[, dim])`` each; the frame as ColumnBatcher writes it.""" + cols, body = [], b"" + for name, ctype, null_rows, data, *dim in columns: + nulls = _bits(count, null_rows) + head = {"name": name, "type": ctype, "nulls": len(nulls), "bytes": len(data)} + if dim: + head["dim"] = dim[0] + cols.append(head) + body += nulls + data + header = json.dumps({"count": count, "cols": cols}).encode("utf-8") + return memoryview(struct.pack(" 1 or t.num_rows == 8000 + + +def test_to_columns_of_the_stored_rows(table): + db, rows = table + c = db.query("sql", _sel()).to_columns() + assert c["big"].dtype == np.int64 and list(c["big"]) == [r["big"] for r in rows] + assert c["n"].dtype == np.float64 # an int column with nulls + assert c["d"].dtype == np.float64 and c["ts"].dtype == np.dtype("datetime64[ms]") + assert c["v9"].dtype == np.float32 and c["v9"].shape == (5000, 9) + assert c["v128"].shape == (5000, 128) + np.testing.assert_array_equal( + c["v9"], np.array([r["v9"] for r in rows], dtype=np.float32) + ) + assert c["s"] == [r["s"] for r in rows] + assert c["ok"] == [r["ok"] for r in rows] # a bool column with nulls is a list + + +def test_column_subset_and_order(table): + db, rows = table + t = db.query("sql", _sel()).to_arrow(columns=["s", "big"]) + assert t.column_names == ["s", "big"] + assert t.column("big").to_pylist() == [r["big"] for r in rows] + c = db.query("sql", _sel()).to_columns(columns=["v3"]) + assert list(c) == ["v3"] and c["v3"].shape == (5000, 3) + + +def test_empty_and_one_row_results(table): + db, rows = table + empty = db.query("sql", "SELECT n, s FROM T WHERE big < 0").to_arrow() + assert empty.num_rows == 0 + assert db.query("sql", "SELECT n, s FROM T WHERE big < 0").to_columns() == {} + one = db.query( + "sql", "SELECT n, s, v3 FROM T WHERE big = ?", rows[1]["big"] + ).to_arrow() + assert one.num_rows == 1 + assert one.column("s").to_pylist() == [rows[1]["s"]] + assert one.column("v3").to_pylist() == [rows[1]["v3"]] + assert db.query("sql", "SELECT FROM T WHERE big < 0").to_arrow().num_rows == 0 + + +def test_a_cursor_read_by_row_first_then_to_arrow(table): + db, rows = table + rs = db.query("sql", _sel()) + first = [next(rs) for _ in range(10)] + assert [r.get("big") for r in first] == [r["big"] for r in rows[:10]] + t = rs.to_arrow() + assert t.column("big").to_pylist() == [r["big"] for r in rows[10:]] + + +def test_many_columns(temp_db): + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE Many") + names = [f"c{k}" for k in range(80)] + for k, name in enumerate(names): + db.command( + "sql", + f"CREATE PROPERTY Many.{name} {['INTEGER', 'STRING', 'DOUBLE'][k % 3]}", + ) + rows = [ + { + name: (i + k if k % 3 == 0 else f"{i}-{k}" if k % 3 == 1 else (i + k) / 4) + for k, name in enumerate(names) + } + for i in range(600) + ] + db.insert_many("Many", rows) + t = db.query("sql", "SELECT FROM Many ORDER BY c0").to_arrow(columns=names) + assert t.column_names == names + for name in ("c0", "c1", "c2", "c79"): + assert t.column(name).to_pylist() == [r[name] for r in rows] + + +def test_to_dataframe_and_to_polars(table): + db, rows = table + df = db.query("sql", _sel()).to_dataframe() + assert list(df.columns) == [n for n, _ in TYPES] + assert list(df["big"]) == [r["big"] for r in rows] + assert df["s"].tolist() == [r["s"] for r in rows] + assert [list(v) for v in df["v3"]] == [r["v3"] for r in rows] + assert df["ts"].iloc[1] == datetime(2026, 1, 2, 3, 4, 5, 1000) + import polars as pl + + p = db.query("sql", _sel()).to_polars() + assert isinstance(p, pl.DataFrame) and p.height == 5000 + assert p["big"].to_list() == [r["big"] for r in rows] diff --git a/tests/test_native_column_frames_reference.py b/tests/test_native_column_frames_reference.py new file mode 100644 index 00000000..19b46cab --- /dev/null +++ b/tests/test_native_column_frames_reference.py @@ -0,0 +1,434 @@ +"""The native column frames give the table the JSON row path gave (refs #312). + +Before the frames, ``NativeResultSet.to_arrow``, ``to_columns`` and ``to_dataframe`` read every row as JSON text into one Python dict and +built the columns from those dicts. The module below (``REF_SOURCE``) is that assembly, copied from ``arcadedb._native.columns`` as it was +the day the frames came, and kept here as a private reference: each case reads the same query twice, once through the old route (the +rows drained as dicts, then the reference rules) and once through the frames, and compares the two column by column, value by value. + +Equal output is the bar. Two types are different on purpose, because the JVM wheel's frames type them so and both wheels now read frames +with one decoder: a column of ``float[]`` is a ``fixed_size_list`` (the old rules gave ``list``, the same values: a +float32 widened to float64 is exact), and a column of only nulls is a string column (the old rules gave the null type). + +conftest.py collects this file on the native backend only (the old route is the native build's JSON reply). +""" + +import math +import types +from datetime import date, datetime +from decimal import Decimal + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest + +REF_SOURCE = '''"""Rows to columns for the native backend (the columnar path is Python over the rows; the JPype path decodes a Java buffer).""" + +from datetime import date, datetime, timezone +from decimal import Decimal + + +def _names(rows, columns): + if columns is not None: + return list(columns) + seen = {} + for row in rows: + for k in row: + seen.setdefault(k, None) + return list(seen) + + +def _kind(values): + """'b', 'i', 'f', 'dt', 'dec', 'obj' for the non-null values of a column.""" + kinds = set() + for v in values: + if v is None: + continue + t = type(v) + if t is bool: + kinds.add("b") + elif t is int: + kinds.add("i") + elif t is float: + kinds.add("f") + elif t is datetime or t is date: + kinds.add("dt") + elif t is Decimal: + kinds.add("dec") + else: + kinds.add("obj") + if not kinds: + return "obj" + if len(kinds) == 1: + return kinds.pop() + if kinds == {"i", "f"}: + return "f" + return "obj" + + +def _ms(v): + if isinstance(v, datetime): + if v.tzinfo is not None: + v = v.astimezone(timezone.utc).replace(tzinfo=None) + return v + return datetime(v.year, v.month, v.day) + + +def _float_vectors(values): + """A column of equal-length float lists (None allowed) that float32 holds exactly, as a (count, dim) float32 array with + NaN rows for the nulls; None for any other column. Rows cannot say float[] from List, so only a column whose every + value survives float32 unchanged is read as a vector column (the JPype path decides on the Java element type). + """ + import numpy as np + + dim = 0 + for v in values: + if v is None: + continue + if type(v) is not list or not v or not all(type(x) is float for x in v): + return None + if dim == 0: + dim = len(v) + elif len(v) != dim: + return None + if dim == 0: + return None + arr = np.full((len(values), dim), np.nan, dtype=np.float32) + for i, v in enumerate(values): + if v is not None: + arr[i] = v + if not np.array_equal(arr[i].astype(np.float64), np.array(v)): + return None + return arr + + +def to_columns(rows, columns=None): + import numpy as np + + out = {} + for name in _names(rows, columns): + values = [row.get(name) for row in rows] + if _kind(values) == "obj": + vectors = _float_vectors(values) + if vectors is not None: + out[name] = vectors + continue + has_null = any(v is None for v in values) + kind = _kind(values) + if kind == "i": + out[name] = ( + np.array([np.nan if v is None else v for v in values], dtype=np.float64) + if has_null + else np.array(values, dtype=np.int64) + ) + elif kind == "f": + out[name] = np.array( + [np.nan if v is None else v for v in values], dtype=np.float64 + ) + elif kind == "b": + out[name] = values if has_null else np.array(values, dtype=bool) + elif kind == "dt": + out[name] = np.array( + [ + ( + np.datetime64("NaT", "ms") + if v is None + else np.datetime64(_ms(v), "ms") + ) + for v in values + ], + dtype="datetime64[ms]", + ) + elif kind == "dec": + arr = np.empty(len(values), dtype=object) + arr[:] = values + out[name] = arr + else: + out[name] = values + return out + + +def to_arrow(rows, columns=None): + import pyarrow as pa + + arrays = {} + for name in _names(rows, columns): + values = [row.get(name) for row in rows] + kind = _kind(values) + if kind == "i": + arrays[name] = pa.array(values, type=pa.int64()) + elif kind == "f": + try: + arrays[name] = pa.array(values, type=pa.float64()) + except ( + pa.ArrowInvalid + ): # an int beyond +-2**53 next to floats: the column degrades to exact text + arrays[name] = pa.array( + [None if v is None else str(v) for v in values], type=pa.string() + ) + elif kind == "b": + arrays[name] = pa.array(values, type=pa.bool_()) + elif kind == "dt": + arrays[name] = pa.array( + [None if v is None else _ms(v) for v in values], type=pa.timestamp("ms") + ) + elif kind == "dec": + try: + arrays[name] = pa.array(values) + except ( + pa.ArrowInvalid + ): # wider than decimal256 (76 digits): an exact string + arrays[name] = pa.array( + [None if v is None else str(v) for v in values], type=pa.string() + ) + else: + try: + arrays[name] = pa.array([None if v is None else v for v in values]) + except (pa.ArrowInvalid, pa.ArrowTypeError): + arrays[name] = pa.array( + [None if v is None else str(v) for v in values], type=pa.string() + ) + return pa.table(arrays) +''' + +REF = types.ModuleType("columns_reference") +exec( + compile(REF_SOURCE, "columns_reference", "exec"), REF.__dict__ +) # nosec B102 - the test's own copy of the old rules + + +def _f32(x): + return float(np.float32(x)) + + +def _row(i): + null = i % 7 == 3 + row = { + "n": None if null else i, + "big": (1 << 40) + i, + "sh": None if null else (i % 2000) - 1000, + "by": None if null else (i % 200) - 100, + "d": i / 3, + "f": None if null else _f32(i / 7), + "s": None if null else f"row {i} \U0001f600 \x00 \u00e9", + "ok": None if null else i % 2 == 0, + "ts": datetime(2026, 1, 2, 3, 4, 5, (i % 1000) * 1000), + "day": date(2000 + i % 30, 1 + i % 12, 1 + i % 28), + "m": Decimal(f"{i}.25"), + "v1": [_f32(i)], + "v3": [_f32(i), _f32(i / 2), 0.5], + "v9": None if null else [_f32(i + k / 10) for k in range(9)], + "v128": [_f32(i * 0.5 + k) for k in range(128)], + "pad": "p" * 400, + "nothing": None, + } + return row + + +TYPES = [ + ("n", "INTEGER"), + ("big", "LONG"), + ("sh", "SHORT"), + ("by", "BYTE"), + ("d", "DOUBLE"), + ("f", "FLOAT"), + ("s", "STRING"), + ("ok", "BOOLEAN"), + ("ts", "DATETIME"), + ("day", "DATE"), + ("m", "DECIMAL"), + ("v1", "ARRAY_OF_FLOATS"), + ("v3", "ARRAY_OF_FLOATS"), + ("v9", "ARRAY_OF_FLOATS"), + ("v128", "ARRAY_OF_FLOATS"), + ("pad", "STRING"), + ("nothing", "STRING"), +] +NAMES = [n for n, _ in TYPES] +ROWS = 5000 + + +@pytest.fixture(scope="module") +def ref_db(tmp_path_factory): + import arcadedb + + path = str(tmp_path_factory.mktemp("frames_ref") / "db") + db = arcadedb.create_database(path) + db.command("sql", "CREATE DOCUMENT TYPE R") + for name, type_name in TYPES: + db.command("sql", f"CREATE PROPERTY R.{name} {type_name}") + db.insert_many("R", [_row(i) for i in range(ROWS)]) + yield db + db.close() + + +def _norm(v): + """A value as a plain Python object, NaN, NaT and None alike (pandas turns a None into NaN in some columns).""" + if isinstance(v, np.ndarray): + return [_norm(x) for x in v.tolist()] + if isinstance(v, (list, tuple)): + return [_norm(x) for x in v] + if v is None or v is pd.NaT: + return None + if isinstance(v, float) and math.isnan(v): + return None + if isinstance(v, np.datetime64): + return None if np.isnat(v) else v.astype("datetime64[us]").item() + if isinstance(v, np.generic): + return v.item() + if isinstance(v, pd.Timestamp): + return v.to_pydatetime() + return v + + +def _old_rows(db, sql, *args): + return db.query("sql", sql, *args)._drain_dicts() + + +def _assert_arrow_equal(new, old): + assert new.column_names == old.column_names + assert new.num_rows == old.num_rows + for name in old.column_names: + have, want = new.column(name), old.column(name) + if pa.types.is_null(want.type): + assert have.type == pa.string() and have.null_count == len(have), name + elif pa.types.is_list(want.type) and want.type.value_type == pa.float64(): + assert pa.types.is_fixed_size_list(have.type) or pa.types.is_list( + have.type + ), name + assert have.type.value_type in (pa.float32(), pa.float64()), name + else: + assert have.type == want.type, (name, have.type, want.type) + assert [_norm(x) for x in have.to_pylist()] == [ + _norm(x) for x in want.to_pylist() + ], name + + +def _assert_columns_equal(new, old): + assert list(new) == list(old) + for name, want in old.items(): + have = new[name] + if isinstance(want, np.ndarray): + assert isinstance(have, np.ndarray), (name, type(have)) + assert have.dtype == want.dtype, (name, have.dtype, want.dtype) + assert have.shape == want.shape, name + else: + assert isinstance(have, list), (name, type(have)) + assert _norm(have) == _norm(want), name + + +def _assert_frame_equal(new, old_rows, columns=None): + old = pd.DataFrame(old_rows) + if columns is not None: + old = old.reindex(columns=list(columns)) + assert list(new.columns) == list(old.columns) + for name in old.columns: + assert _norm(new[name].tolist()) == _norm(old[name].tolist()), name + + +QUERIES = [ + ("full", "SELECT FROM R ORDER BY big", None), + ( + "projection", + "SELECT n, big, sh, by, d, f, s, ok, ts, day, m FROM R ORDER BY big", + None, + ), + ("vectors", "SELECT big, v1, v3, v9, v128 FROM R ORDER BY big", None), + ("limit 1", "SELECT FROM R ORDER BY big LIMIT 1", None), + ("limit 100, one reply", "SELECT FROM R ORDER BY big LIMIT 100", None), + ("limit 300, a reply and a cursor", "SELECT FROM R ORDER BY big LIMIT 300", None), + ("empty", "SELECT FROM R WHERE big < 0", None), + ("subset", "SELECT FROM R ORDER BY big", ["s", "big", "v3", "ts"]), + ("subset with a column nobody has", "SELECT FROM R ORDER BY big", ["n", "absent"]), + ("subset on a short result", "SELECT FROM R ORDER BY big LIMIT 20", ["m", "ok"]), +] + + +@pytest.mark.parametrize("label,sql,columns", QUERIES, ids=[q[0] for q in QUERIES]) +def test_to_arrow_equals_the_old_route(ref_db, label, sql, columns): + old = REF.to_arrow(_old_rows(ref_db, sql), columns) + new = ref_db.query("sql", sql).to_arrow(columns=columns) + _assert_arrow_equal(new, old) + + +@pytest.mark.parametrize("label,sql,columns", QUERIES, ids=[q[0] for q in QUERIES]) +def test_to_columns_equals_the_old_route(ref_db, label, sql, columns): + old = REF.to_columns(_old_rows(ref_db, sql), columns) + new = ref_db.query("sql", sql).to_columns(columns=columns) + _assert_columns_equal(new, old) + + +@pytest.mark.parametrize( + "label,sql,columns", QUERIES[:7], ids=[q[0] for q in QUERIES[:7]] +) +def test_to_dataframe_equals_the_old_route(ref_db, label, sql, columns): + rows = _old_rows(ref_db, sql) + new = ref_db.query("sql", sql).to_dataframe() + _assert_frame_equal(new, rows) + + +@pytest.mark.parametrize("batch_size", [1, 64, 255, 256, 257, 1000, 25000]) +def test_the_frames_do_not_depend_on_the_batch_size(ref_db, batch_size): + sql = "SELECT FROM R ORDER BY big" + old = REF.to_arrow(_old_rows(ref_db, sql)) + _assert_arrow_equal(ref_db.query("sql", sql).to_arrow(batch_size=batch_size), old) + + +def test_a_result_larger_than_one_frame_is_read_in_several(ref_db): + from unittest import mock + + from arcadedb._native import lib + + sizes = [] + real = lib.call_frame + + def spy(*args): + raw = real(*args) + sizes.append(len(raw)) + return raw + + # 5000 rows of about 2.4 KiB of values: several frames of about 1 MiB each, and the table is the same as one read of the rows + with mock.patch.object(lib, "call_frame", spy): + t = ref_db.query("sql", "SELECT FROM R ORDER BY big").to_arrow() + assert t.num_rows == ROWS + assert len(sizes) >= 4, sizes + assert max(sizes) < 4 * (1 << 20), sizes # a frame, not the table + assert t.column("big").to_pylist() == [(1 << 40) + i for i in range(ROWS)] + + +def test_rows_read_by_iteration_first_then_the_rest_as_columns(ref_db): + sql = "SELECT FROM R ORDER BY big" + old = REF.to_arrow(_old_rows(ref_db, sql)).slice(40) + rs = ref_db.query("sql", sql) + taken = [next(rs).get("big") for _ in range(40)] + assert taken == [(1 << 40) + i for i in range(40)] + _assert_arrow_equal(rs.to_arrow(), old) + + +def test_rows_read_past_the_first_reply_then_the_rest_as_columns(ref_db): + sql = "SELECT FROM R ORDER BY big" + old = REF.to_arrow(_old_rows(ref_db, sql)).slice(300) + rs = ref_db.query("sql", sql) + for _ in range(300): + next(rs) + _assert_arrow_equal(rs.to_arrow(), old) + + +def test_a_result_is_read_once(ref_db): + rs = ref_db.query("sql", "SELECT FROM R ORDER BY big LIMIT 500") + assert rs.to_arrow().num_rows == 500 + assert rs.to_arrow().num_rows == 0 + + +def test_to_list_still_reads_the_json_rows(ref_db): + rows = ref_db.query("sql", "SELECT FROM R ORDER BY big LIMIT 300").to_list() + assert len(rows) == 300 and rows[299]["big"] == (1 << 40) + 299 + + +def test_the_cursor_keeps_no_rows_after_a_row_read(ref_db): + rs = ref_db.query("sql", "SELECT FROM R ORDER BY big") + assert rs._head_n == 256 + for _ in range(300): + next(rs) + assert rs._head_n == 0 + rs.close() From 71bad36ed15678d78e96b62a81d0d03229483738 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Mon, 9 Nov 2026 19:00:08 +0900 Subject: [PATCH 2/6] Tests: the frames test builds its SQL as one literal (bandit B608) Co-Authored-By: Claude Sonnet 5.5 --- tests/test_column_frames.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_column_frames.py b/tests/test_column_frames.py index a2e51a22..0265a17a 100644 --- a/tests/test_column_frames.py +++ b/tests/test_column_frames.py @@ -197,8 +197,8 @@ def table(temp_db): return db, rows -def _sel(extra=""): - return "SELECT n, big, d, f, s, ok, ts, v3, v9, v128 FROM T ORDER BY big" + extra +def _sel(): + return "SELECT n, big, d, f, s, ok, ts, v3, v9, v128 FROM T ORDER BY big" def _same(have, want): From 1ff1e64365a66b9751854a2e1ea736f577b23e3a Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Mon, 9 Nov 2026 20:51:56 +0900 Subject: [PATCH 3/6] Native frames: a frame that ends before the replayed first-reply rows keeps the rest for the next call With batch_size=64 the first frame took 64 of the 256 replayed rows and the cursor then dropped the other 192 (test_to_arrow_multi_batch: 108 rows of 300). The library now puts back the unread tail of the head after each frame. Also: the reference test named a property "by" (an SQL keyword), and the dataframe test expected None where pandas 3 shows nan for a missing string. Co-Authored-By: Claude Sonnet 5.5 --- .../src/main/java/org/humemai/nativelib/ArcadeNative.java | 5 ++++- .../src/main/java/org/humemai/nativelib/ColumnFrames.java | 7 +++++-- tests/test_column_frames.py | 7 ++++++- tests/test_native_column_frames_reference.py | 6 +++--- 4 files changed, 18 insertions(+), 7 deletions(-) diff --git a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java index 52f9d3c5..836e40a7 100644 --- a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java +++ b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java @@ -337,7 +337,7 @@ public static CCharPointer nextBin(final IsolateThread thread, final long rsId, * the negated byte count of the error JSON and that JSON. max: rows at most (a frame also ends at about 1 MiB of values). columns: a * JSON array of the column names to read, or "" for the union of the names of the rows. flags bit 0: datetimes as microseconds (the * to_dataframe form). skip: of the rows that the first reply of this cursor already held (HEADS), the number the caller has taken; the - * first frame starts after them. The frame that ends the rows closes the cursor. The reply is freed with arcadedb_free. + * first frame starts after them, and a frame that ends before those rows do leaves the rest in HEADS for the next call (skip 0). The frame that ends the rows closes the cursor. The reply is freed with arcadedb_free. */ @CEntryPoint(name = "arcadedb_n_next_col") @SuppressWarnings("unchecked") @@ -361,6 +361,9 @@ public static CCharPointer nextCol(final IsolateThread thread, final long rsId, if (frame.done) { removeCursor(rsId); closeQuietly(rs); + } else if (head != null && frame.headUsed < head.size()) { + // the frame ended (batch_size rows or about 1 MiB) before the replayed rows did: the next call starts with the rest + HEADS.put(rsId, new ArrayList<>(head.subList(frame.headUsed, head.size()))); } return frameOut(frame.bytes.length, frame.bytes); } catch (final Throwable t) { diff --git a/native/src/main/java/org/humemai/nativelib/ColumnFrames.java b/native/src/main/java/org/humemai/nativelib/ColumnFrames.java index 0c507346..a45aad38 100644 --- a/native/src/main/java/org/humemai/nativelib/ColumnFrames.java +++ b/native/src/main/java/org/humemai/nativelib/ColumnFrames.java @@ -57,10 +57,13 @@ private ColumnFrames() { static final class Frame { final byte[] bytes; final boolean done; + /** The index in {@code head} of the first row this frame did not take (a frame may end before the head does). */ + final int headUsed; - Frame(final byte[] bytes, final boolean done) { + Frame(final byte[] bytes, final boolean done, final int headUsed) { this.bytes = bytes; this.done = done; + this.headUsed = headUsed; } } @@ -377,6 +380,6 @@ else if (type.equals("dtu")) out.putInt(headerBytes.length); out.put(headerBytes); out.put(bodyBytes); - return new Frame(out.array(), done); + return new Frame(out.array(), done, h); } } diff --git a/tests/test_column_frames.py b/tests/test_column_frames.py index 0265a17a..7956b15a 100644 --- a/tests/test_column_frames.py +++ b/tests/test_column_frames.py @@ -14,6 +14,7 @@ from datetime import datetime import numpy as np +import pandas as pd import pyarrow as pa import pytest @@ -328,7 +329,11 @@ def test_to_dataframe_and_to_polars(table): df = db.query("sql", _sel()).to_dataframe() assert list(df.columns) == [n for n, _ in TYPES] assert list(df["big"]) == [r["big"] for r in rows] - assert df["s"].tolist() == [r["s"] for r in rows] + # pandas decides how a missing string is shown (None in an object column, nan in the string dtype of pandas 3): the same on + # both wheels, since both build the frame from a list with None + assert [None if pd.isna(v) else v for v in df["s"].tolist()] == [ + r["s"] for r in rows + ] assert [list(v) for v in df["v3"]] == [r["v3"] for r in rows] assert df["ts"].iloc[1] == datetime(2026, 1, 2, 3, 4, 5, 1000) import polars as pl diff --git a/tests/test_native_column_frames_reference.py b/tests/test_native_column_frames_reference.py index 19b46cab..630a4390 100644 --- a/tests/test_native_column_frames_reference.py +++ b/tests/test_native_column_frames_reference.py @@ -207,7 +207,7 @@ def _row(i): "n": None if null else i, "big": (1 << 40) + i, "sh": None if null else (i % 2000) - 1000, - "by": None if null else (i % 200) - 100, + "bt": None if null else (i % 200) - 100, "d": i / 3, "f": None if null else _f32(i / 7), "s": None if null else f"row {i} \U0001f600 \x00 \u00e9", @@ -229,7 +229,7 @@ def _row(i): ("n", "INTEGER"), ("big", "LONG"), ("sh", "SHORT"), - ("by", "BYTE"), + ("bt", "BYTE"), ("d", "DOUBLE"), ("f", "FLOAT"), ("s", "STRING"), @@ -330,7 +330,7 @@ def _assert_frame_equal(new, old_rows, columns=None): ("full", "SELECT FROM R ORDER BY big", None), ( "projection", - "SELECT n, big, sh, by, d, f, s, ok, ts, day, m FROM R ORDER BY big", + "SELECT n, big, sh, bt, d, f, s, ok, ts, day, m FROM R ORDER BY big", None, ), ("vectors", "SELECT big, v1, v3, v9, v128 FROM R ORDER BY big", None), From d606ccee42bacd57f0bde43fbc399879a023158c Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Mon, 9 Nov 2026 21:31:08 +0900 Subject: [PATCH 4/6] Native frames, second round: no replay of the first reply, types that belong to the whole result, bounded frames The first reply of a query stays what the caller already holds as dicts and is written as one more frame by Python (colframes.py) by the rules of the library's encoder, so the cursor keeps no rows (the replay kept up to 256 live Result objects per open cursor and read them after the snapshot the JSON reply had taken). The JSON reply names the properties that were float[] or double[] (fv, dv): a vector column is a declared array, not a list that happens to fit float32, at any result size. columns.py, with its size-dependent guess, is gone. Frames: the row limit is clamped to a C int (a batch_size of 2**31 or more wrapped to zero and looped), a frame ends by the larger of the values read and rows x columns x 9 bytes (a schemaless result with a property per row allocated gigabytes against a 1 MiB estimate), an empty frame from an open cursor is an error, a NULL reply and a library older than the wheel are errors with a message. Types: float32 next to float64 or another dimension unify to a list of float (never text), a frame of only nulls takes the type of the others, an integer beyond 2**53 next to a float stays exact in a json column on both wheels (ColumnBatcher too), an empty result with columns= returns those columns on both wheels, to_columns arrays are writable, to_dataframe vectors are numpy float32 cells and datetimes datetime64[us] on both wheels (documented). Co-Authored-By: Claude Sonnet 5.5 --- .../org/humemai/nativelib/ArcadeNative.java | 55 ++-- .../org/humemai/nativelib/ColumnFrames.java | 55 ++-- .../org/humemai/nativelib/NativeCodec.java | 47 ++- src/arcadedb/_columnar.py | 82 +++++- src/arcadedb/_native/colframes.py | 269 ++++++++++++------ src/arcadedb/_native/columns.py | 167 ----------- src/arcadedb/_native/database.py | 77 +++-- src/arcadedb/_native/lib.py | 18 +- src/arcadedb/results.py | 49 ++-- .../com/arcadedb/python/ColumnBatcher.java | 15 +- tests/test_column_frames.py | 190 +++++++++++++ tests/test_native_column_frames_reference.py | 159 ++++++++--- 12 files changed, 751 insertions(+), 432 deletions(-) delete mode 100644 src/arcadedb/_native/columns.py diff --git a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java index 836e40a7..13ce33a2 100644 --- a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java +++ b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java @@ -20,7 +20,6 @@ import com.arcadedb.engine.ComponentFile; import com.arcadedb.exception.DuplicatedKeyException; import com.arcadedb.exception.NeedRetryException; -import com.arcadedb.query.sql.executor.Result; import com.arcadedb.query.sql.executor.ResultSet; import com.arcadedb.schema.VertexType; import com.arcadedb.server.ArcadeDBServer; @@ -37,7 +36,6 @@ import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; -import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -47,17 +45,6 @@ public final class ArcadeNative { static final Map DBS = new ConcurrentHashMap<>(); private static final Map RESULT_SETS = new ConcurrentHashMap<>(); - /** - * The rows of the first reply of an open cursor, kept until the cursor is read again or closed: that reply is JSON for the row APIs, - * and a columnar read (arcadedb_n_next_col) that starts from it replays these rows into its first frame instead of losing them. - * At most the first reply (256 rows) of a cursor that has more rows. - */ - private static final Map> HEADS = new ConcurrentHashMap<>(); - - private static ResultSet removeCursor(final long id) { - HEADS.remove(id); - return RESULT_SETS.remove(id); - } private static final Map SERVERS = new ConcurrentHashMap<>(); static final AtomicLong IDS = new AtomicLong(1); @@ -301,10 +288,9 @@ public static CCharPointer next(final IsolateThread thread, final long rsId, fin rs = RESULT_SETS.get(rsId); if (rs == null) throw new IllegalArgumentException("unknown result set " + rsId); - HEADS.remove(rsId); - return out(batchOf(rsId, rs, max, null)); + return out(batchOf(rsId, rs, max)); } catch (final Throwable t) { - removeCursor(rsId); + RESULT_SETS.remove(rsId); closeQuietly(rs); return err(t); } @@ -323,10 +309,9 @@ public static CCharPointer nextBin(final IsolateThread thread, final long rsId, if (rs == null) throw new IllegalArgumentException("unknown result set " + rsId); final ResultSet cursor = rs; - HEADS.remove(rsId); - return withFloatTail(() -> batchOf(rsId, cursor, max, null)); + return withFloatTail(() -> batchOf(rsId, cursor, max)); } catch (final Throwable t) { - removeCursor(rsId); + RESULT_SETS.remove(rsId); closeQuietly(rs); return err(t); } @@ -336,19 +321,18 @@ public static CCharPointer nextBin(final IsolateThread thread, final long rsId, * The next rows of an open cursor as one column frame (see ColumnFrames), written as an int64 LE byte count and the frame; a failure is * the negated byte count of the error JSON and that JSON. max: rows at most (a frame also ends at about 1 MiB of values). columns: a * JSON array of the column names to read, or "" for the union of the names of the rows. flags bit 0: datetimes as microseconds (the - * to_dataframe form). skip: of the rows that the first reply of this cursor already held (HEADS), the number the caller has taken; the - * first frame starts after them, and a frame that ends before those rows do leaves the rest in HEADS for the next call (skip 0). The frame that ends the rows closes the cursor. The reply is freed with arcadedb_free. + * to_dataframe form). The rows the cursor already gave out in a JSON reply are not in the frame. The frame that ends the rows closes the + * cursor. The reply is freed with arcadedb_free. */ @CEntryPoint(name = "arcadedb_n_next_col") @SuppressWarnings("unchecked") public static CCharPointer nextCol(final IsolateThread thread, final long rsId, final int max, final CCharPointer columns, - final int flags, final int skip) { + final int flags) { ResultSet rs = null; try { rs = RESULT_SETS.get(rsId); if (rs == null) throw new IllegalArgumentException("unknown result set " + rsId); - final List head = HEADS.remove(rsId); String[] pinned = null; final String spec = str(columns); if (spec != null && !spec.isEmpty()) { @@ -357,17 +341,14 @@ public static CCharPointer nextCol(final IsolateThread thread, final long rsId, for (int i = 0; i < pinned.length; i++) pinned[i] = (String) names.get(i); } - final ColumnFrames.Frame frame = ColumnFrames.next(rs, head, skip, max, pinned, (flags & 1) != 0); + final ColumnFrames.Frame frame = ColumnFrames.next(rs, max, pinned, (flags & 1) != 0); if (frame.done) { - removeCursor(rsId); + RESULT_SETS.remove(rsId); closeQuietly(rs); - } else if (head != null && frame.headUsed < head.size()) { - // the frame ended (batch_size rows or about 1 MiB) before the replayed rows did: the next call starts with the rest - HEADS.put(rsId, new ArrayList<>(head.subList(frame.headUsed, head.size()))); } return frameOut(frame.bytes.length, frame.bytes); } catch (final Throwable t) { - removeCursor(rsId); + RESULT_SETS.remove(rsId); closeQuietly(rs); final byte[] e = errorJson(t).getBytes(StandardCharsets.UTF_8); return frameOut(-e.length, e); @@ -384,7 +365,7 @@ private static CCharPointer frameOut(final int length, final byte[] b) { @CEntryPoint(name = "arcadedb_n_rs_close") public static CCharPointer rsClose(final IsolateThread thread, final long rsId) { - closeQuietly(removeCursor(rsId)); + closeQuietly(RESULT_SETS.remove(rsId)); return OK(); } @@ -393,20 +374,16 @@ private static String batch(final ResultSet rs, final int max) { return "{\"ok\":true,\"rs\":0,\"rows\":[],\"meta\":[],\"drained\":true}"; final long id = IDS.getAndIncrement(); RESULT_SETS.put(id, rs); - final List keep = new ArrayList<>(); - final String reply = batchOf(id, rs, max, keep); - if (RESULT_SETS.containsKey(id)) - HEADS.put(id, keep); - return reply; + return batchOf(id, rs, max); } - private static String batchOf(final long id, final ResultSet rs, final int max, final List keep) { + private static String batchOf(final long id, final ResultSet rs, final int max) { final StringBuilder sb = new StringBuilder("{\"ok\":true,\"rs\":"); final int mark = sb.length(); sb.append(id).append(','); - final boolean drained = NativeCodec.appendBatch(sb, rs, max, keep); + final boolean drained = NativeCodec.appendBatch(sb, rs, max); if (drained) { - removeCursor(id); + RESULT_SETS.remove(id); closeQuietly(rs); sb.replace(mark, mark + String.valueOf(id).length(), "0"); } @@ -719,7 +696,7 @@ public static CCharPointer lastError(final IsolateThread thread) { public static void rsCloseMany(final IsolateThread thread, final CLongPointer ids, final int n) { try { for (int i = 0; i < n; i++) - closeQuietly(removeCursor(ids.read(i))); + closeQuietly(RESULT_SETS.remove(ids.read(i))); } catch (final Throwable ignored) { // closing is best effort, as in arcadedb_n_rs_close } diff --git a/native/src/main/java/org/humemai/nativelib/ColumnFrames.java b/native/src/main/java/org/humemai/nativelib/ColumnFrames.java index a45aad38..02d9106c 100644 --- a/native/src/main/java/org/humemai/nativelib/ColumnFrames.java +++ b/native/src/main/java/org/humemai/nativelib/ColumnFrames.java @@ -57,13 +57,10 @@ private ColumnFrames() { static final class Frame { final byte[] bytes; final boolean done; - /** The index in {@code head} of the first row this frame did not take (a frame may end before the head does). */ - final int headUsed; - Frame(final byte[] bytes, final boolean done, final int headUsed) { + Frame(final byte[] bytes, final boolean done) { this.bytes = bytes; this.done = done; - this.headUsed = headUsed; } } @@ -126,38 +123,33 @@ else if (v instanceof Map m) } /** - * The next frame. {@code head} holds rows that were already read from the cursor when its first reply was encoded as JSON (the - * frame starts at {@code head[skip]}, then goes on with the cursor), or is null. {@code pinned} is the column set, or null to take - * the union of the names of the rows of this frame, in order of first appearance. With {@code micros} a naive datetime column is - * epoch microseconds ("dtu") and a column with a DATE, a time-zone-aware temporal or a year Python cannot hold is "json". + * The next frame of the cursor. {@code pinned} is the column set, or null to take the union of the names of the rows of this frame, in + * order of first appearance. With {@code micros} a naive datetime column is epoch microseconds ("dtu") and a column with a DATE, a + * time-zone-aware temporal or a year Python cannot hold is "json". + * + * The frame ends at {@code max} rows or when what it will hold reaches FRAME_BYTES: the larger of the sizes of the values read and of the + * fixed-width buffers (rows times columns times 9 bytes, which a schemaless result with a different property in every row makes much + * bigger than its values), so the allocation of a frame stays near its budget whatever the shape of the rows. A frame has one row at least. */ - static Frame next(final ResultSet rs, final List head, final int skip, final int max, final String[] pinned, - final boolean micros) { - final List results = new ArrayList<>(Math.min(Math.max(max, 1), 1 << 14)); - int h = Math.max(skip, 0); + static Frame next(final ResultSet rs, final int max, final String[] pinned, final boolean micros) { + final int cap = Math.max(max, 1); + final List results = new ArrayList<>(Math.min(cap, 1 << 14)); + final Set names = new LinkedHashSet<>(); + String[] previous = null; long budget = 0; boolean done = false; - while (results.size() < max && budget < FRAME_BYTES) { - final Result row; - if (head != null && h < head.size()) - row = head.get(h++); - else if (rs.hasNext()) - row = rs.next(); - else { + while (results.size() < cap + && Math.max(budget, (long) results.size() * (pinned != null ? pinned.length : names.size()) * 9L) < FRAME_BYTES) { + if (!rs.hasNext()) { done = true; break; } + final Result row = rs.next(); results.add(row); budget += estimate(row); - } - if (!done) - done = (head == null || h >= head.size()) && !rs.hasNext(); - - String[] columns = pinned; - if (columns == null) { - final Set names = new LinkedHashSet<>(); - String[] previous = null; - for (final Result row : results) { + if (pinned == null) { + // rows of one type almost always list the same names in the same order: a row is compared with the one before it name by name + // and only a row that differs is hashed into the set final Set rowNames = row.getPropertyNames(); if (previous != null && rowNames.size() == previous.length) { int i = 0; @@ -172,8 +164,11 @@ else if (rs.hasNext()) names.addAll(rowNames); previous = rowNames.toArray(new String[0]); } - columns = names.toArray(new String[0]); } + if (!done) + done = !rs.hasNext(); + + final String[] columns = pinned != null ? pinned : names.toArray(new String[0]); final int n = columns.length; final int count = results.size(); final List rows = new ArrayList<>(count); @@ -380,6 +375,6 @@ else if (type.equals("dtu")) out.putInt(headerBytes.length); out.put(headerBytes); out.put(bodyBytes); - return new Frame(out.array(), done, h); + return new Frame(out.array(), done); } } diff --git a/native/src/main/java/org/humemai/nativelib/NativeCodec.java b/native/src/main/java/org/humemai/nativelib/NativeCodec.java index 5d037a50..b3bbab06 100644 --- a/native/src/main/java/org/humemai/nativelib/NativeCodec.java +++ b/native/src/main/java/org/humemai/nativelib/NativeCodec.java @@ -51,13 +51,38 @@ private NativeCodec() { * Appends {@code "rows":[..],"meta":[..],"drained":b} (no braces) with up to max rows. meta has one entry per row: 0 * for a projection, or [rid, typeName, kind] for a record row (kind v, e or d). */ - static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int max) { - return appendBatch(sb, rs, max, null); + /** + * The names of the properties whose value was a float[] (f) or a double[] (d) in a row of a batch. A JSON reply writes both as a list of + * numbers, so the reader cannot tell a vector from a list of doubles: the reply names the vectors, {@code "fv":[..]} and {@code "dv":[..]} + * after "drained" (absent when there is none), and a columnar read types those columns as vectors, as a column frame does. + */ + static final class VecHints { + final java.util.LinkedHashSet f = new java.util.LinkedHashSet<>(); + final java.util.LinkedHashSet d = new java.util.LinkedHashSet<>(); + + void appendTo(final StringBuilder sb) { + appendNames(sb, ",\"fv\":", f); + appendNames(sb, ",\"dv\":", d); + } + + private static void appendNames(final StringBuilder sb, final String key, final java.util.Set names) { + if (names.isEmpty()) + return; + sb.append(key).append('['); + boolean first = true; + for (final String name : names) { + if (!first) + sb.append(','); + first = false; + quote(sb, name); + } + sb.append(']'); + } } - /** As above; {@code keep}, when not null, receives each row that was encoded (a columnar read replays them, see ArcadeNative.HEADS). */ - static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int max, final java.util.List keep) { + static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int max) { final StringBuilder meta = new StringBuilder("["); + final VecHints hints = new VecHints(); sb.append("\"rows\":["); int n = 0; boolean drained = false; @@ -76,15 +101,13 @@ static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int sb.append(','); meta.append(','); } - appendRow(sb, meta, row); + appendRow(sb, meta, row, hints); } catch (final Throwable t) { sb.setLength(rowStart); meta.setLength(metaStart); throw t; } n++; - if (keep != null) - keep.add(row); if (sb.length() >= MAX_BATCH_CHARS) break; } @@ -97,10 +120,11 @@ static boolean appendBatch(final StringBuilder sb, final ResultSet rs, final int sb.append("],\"meta\":").append(meta).append("],\"drained\":").append(drained); if (readError != null) sb.append(",\"rerr\":").append(readError); + hints.appendTo(sb); return drained; } - private static void appendRow(final StringBuilder sb, final StringBuilder meta, final Result row) { + private static void appendRow(final StringBuilder sb, final StringBuilder meta, final Result row, final VecHints hints) { sb.append('{'); boolean first = true; for (final String name : row.getPropertyNames()) { @@ -109,7 +133,12 @@ private static void appendRow(final StringBuilder sb, final StringBuilder meta, first = false; quote(sb, name); sb.append(':'); - append(sb, row.getProperty(name)); + final Object value = row.getProperty(name); + if (value instanceof float[]) + hints.f.add(name); + else if (value instanceof double[]) + hints.d.add(name); + append(sb, value); } sb.append('}'); if (row.isElement()) { diff --git a/src/arcadedb/_columnar.py b/src/arcadedb/_columnar.py index 7cb62e36..16c66d81 100644 --- a/src/arcadedb/_columnar.py +++ b/src/arcadedb/_columnar.py @@ -165,6 +165,25 @@ def _unify_arrow_chunk_types(arrs: list, pa) -> list: except pa.ArrowException: pass + if all(pa.types.is_list(t) or pa.types.is_fixed_size_list(t) for t in types): + # vectors and lists of numbers: a dimension that changes between frames makes a list, float32 next to float64 is float64 + # (the widening is exact), never text + values = {t.value_type for t in types} + if values <= {pa.float32(), pa.float64()}: + value = pa.float64() if pa.float64() in values else pa.float32() + sizes = { + t.list_size if pa.types.is_fixed_size_list(t) else None for t in types + } + target = ( + pa.list_(value, next(iter(sizes))) + if len(sizes) == 1 and None not in sizes + else pa.list_(value) + ) + try: + return [pa.array(a.to_pylist(), type=target) for a in arrs] + except pa.ArrowException: + pass + if all(pa.types.is_decimal(t) for t in types): common = _common_decimal_type(list(types), pa) if common is not None: @@ -173,6 +192,31 @@ def _unify_arrow_chunk_types(arrs: list, pa) -> list: return _cast_all_to_string(arrs, pa) +MAX_FRAME_ROWS = 2**31 - 1 # a frame's row limit crosses to the library as a C int + + +def frame_rows(batch_size) -> int: + """The row limit of a frame for ``batch_size``: at least 1 (a ValueError otherwise), at most what a C int holds, so a huge + ``batch_size`` means "as many as one frame takes" instead of wrapping around to zero or a negative number. + """ + rows = int(batch_size) + if rows < 1: + raise ValueError(f"batch_size must be at least 1, got {batch_size}") + return min(rows, MAX_FRAME_ROWS) + + +def empty_columns(columns) -> Dict[str, Any]: + """The ``to_columns`` dict of a result with no rows: no column, or the ones the caller asked for, empty.""" + return {} if columns is None else {name: [] for name in columns} + + +def empty_arrow(columns): + """The ``to_arrow`` table of a result with no rows: no column, or the ones the caller asked for, empty and of the null type.""" + import pyarrow as pa + + return pa.table({} if columns is None else {name: pa.nulls(0) for name in columns}) + + def _null_mask(buf, pos: int, nulls_len: int, count: int): """``(null_bits, mask)`` of a column: ``mask`` is a bool array (True = null) or None when no row is null.""" import numpy as np @@ -213,14 +257,15 @@ def decode_columns_frame( pos += col["bytes"] if ctype == "i8": - arr = np.frombuffer(data, dtype=" bool: + return ( + isinstance(part, list) and len(part) == count and all(v is None for v in part) + ) + + def merge_columns( - names: List[str], batches: List[Tuple[int, Dict[str, Any]]] + names: List[str], batches: List[Tuple[int, Dict[str, Any]]], columns=None ) -> Dict[str, Any]: - """The ``to_columns`` dict of the parts of every frame: a column a frame lacks gets nulls (NaN, NaT or None) there.""" + """The ``to_columns`` dict of the parts of every frame. The type of a column belongs to the whole result, not to where the frames + end: a frame that lacks the column, or holds only nulls for it, takes the type of the frames that have values (NaN, NaT or None + in its rows). With no rows at all the result is :func:`empty_columns`.""" import numpy as np + if not batches or not names: + return empty_columns(columns) out: Dict[str, Any] = {} for name in names: - template = next(p[name] for _, p in batches if name in p) + typed = [ + p[name] for c, p in batches if name in p and not _is_all_none(p[name], c) + ] + template = typed[0] if typed else next(p[name] for _, p in batches if name in p) parts = [ - p[name] if name in p else _missing_column_part(c, template) + ( + p[name] + if name in p and not (typed and _is_all_none(p[name], c)) + else _missing_column_part(c, template) + ) for c, p in batches ] np_parts = [p for p in parts if not isinstance(p, list)] @@ -374,13 +436,15 @@ def decode_arrow_frame( return count, chunks, header -def merge_arrow(names: List[str], batches: List[Tuple[int, Dict[str, Any]]]): +def merge_arrow( + names: List[str], batches: List[Tuple[int, Dict[str, Any]]], columns=None +): """The ``pyarrow.Table`` of the chunks of every frame: one chunked column per name, one type per column (see :func:`_unify_arrow_chunk_types`), a frame that lacks a column adds nulls.""" import pyarrow as pa if not batches or not names: - return pa.table({}) + return empty_arrow(columns) return pa.table( { n: pa.chunked_array( diff --git a/src/arcadedb/_native/colframes.py b/src/arcadedb/_native/colframes.py index e15dcb2f..31e47cc0 100644 --- a/src/arcadedb/_native/colframes.py +++ b/src/arcadedb/_native/colframes.py @@ -1,29 +1,186 @@ """Column frames on the native build: the glue between ``arcadedb_n_next_col`` and the decoder both wheels share (``arcadedb._columnar``). -A result of the native build reaches Python in two forms. The rows after the first reply are column frames, read by the shared decoder. -The first reply of a query (up to ``_FIRST_BATCH`` rows) is JSON text that the row APIs read, and it is the only part of a result that a -columnar read holds as dicts: the rows that are still buffered when the columnar read starts go through the rules of :mod:`.columns` -and are merged with the frames as one more chunk. When the cursor still holds those rows (``arcadedb_n_next_col`` replays them), they -are not taken from the dicts at all; a result that the first reply held whole has no cursor, so its rows are the only form it has. -That chunk goes away with the JSON reply. - -Two things in a chunk of dicts cannot be known from the values: whether a list of floats is a ``float[]`` (a frame says so with its -column type) and nothing else. A column of lists of floats that float32 holds exactly is a vector column in the chunk when the frames -say so, or when there are no frames; when the frames carry it as a list column, the chunk keeps the list. +A result of the native build reaches Python in two forms. The rows after the first reply are column frames, which the library writes and +the shared decoder reads. The first reply of a query (up to ``_FIRST_BATCH`` rows) is JSON text that the row APIs read, and a columnar read +finds those rows (and any the caller read by row) as dicts in the result's buffer. They cannot be read again from the cursor, so they are +written here as one more frame, in the layout the library writes, by the rules the library's encoder follows, and read by the same +decoder: a column has one type however a result splits into the first reply and frames. The JSON reply names the properties that were +``float[]`` or ``double[]`` (``fv`` and ``dv``): a list of numbers in a row cannot say that by itself, and a vector column is a column of +declared arrays, not a column of lists that happen to fit float32. This frame goes away with the JSON reply. """ -from datetime import datetime +import json +import struct +from datetime import date, datetime, timedelta, timezone +from decimal import Decimal from .. import _columnar -from . import columns as _rows from .codec import loads_tagged +_EPOCH = datetime(1970, 1, 1) +_MS = timedelta(milliseconds=1) +_US = timedelta(microseconds=1) +_INT64 = (-(1 << 63), (1 << 63) - 1) +_TWO53 = 1 << 53 + def json_values(data, side, count): - """The values of a ``json`` column of a frame: the native library's tagged JSON, restored as a row's values are.""" + """The values of a ``json`` column of a frame from the library: its tagged JSON, restored as a row's values are.""" return loads_tagged(bytes(data).decode("utf-8"))[0] +def _head_json_values(data, side, count): + """The values of a ``json`` column of the head frame: the buffered values themselves, by the index the column carries.""" + return side[int.from_bytes(bytes(data), "little")] + + +def _naive_utc(v): + if isinstance(v, datetime): + if v.tzinfo is not None: + v = v.astimezone(timezone.utc).replace(tzinfo=None) + return v + return datetime(v.year, v.month, v.day) + + +def _kind(values, micros): + """The column type of a frame that the values would get from the library's encoder: ``i8``, ``f8``, ``b1``, ``dt``, ``dtu``, + ``str``, ``dec`` or ``json``; ``null`` for a column of nothing but nulls.""" + kinds = set() + big_int = False + for v in values: + if v is None: + continue + t = type(v) + if t is bool: + k = "b1" + elif t is int: + if not _INT64[0] <= v <= _INT64[1]: + k = "json" # a BigInteger + else: + k = "i8" + if not -_TWO53 <= v <= _TWO53: + big_int = True + elif t is float: + k = "f8" + elif t is datetime or t is date: + k = "dt" + elif t is str: + k = "str" + elif t is Decimal: + k = "dec" + else: + k = "json" + kinds.add(k) + if k == "json": + break + if not kinds: + return "null" + if kinds == {"i8", "f8"}: + # an integer beyond 2**53 next to a floating value has no exact double: the values stay what they are + return "json" if big_int else "f8" + if len(kinds) > 1: + return "json" + kind = kinds.pop() + if kind == "dt": + if not micros: + return "dt" + if all(v is None or (type(v) is datetime and v.tzinfo is None) for v in values): + return "dtu" + return "json" + return kind + + +def _vector_dim(values): + """The one length of the lists of a column of float lists, or 0 when it is not one (a non-list, an empty or a ragged list).""" + dim = 0 + for v in values: + if v is None: + continue + if type(v) is not list or not v or any(type(x) is not float for x in v): + return 0 + if dim == 0: + dim = len(v) + elif len(v) != dim: + return 0 + return dim + + +def _bitmap(values): + out = bytearray((len(values) + 7) // 8) + for i, v in enumerate(values): + if v is None: + out[i >> 3] |= 1 << (i & 7) + return bytes(out) + + +def _column(values, vector, micros, side): + """``(type, data, dim)`` of one column: the buffer the library's encoder writes for it. ``vector`` is ``"f"`` or ``"d"`` for a + property that was a ``float[]`` or a ``double[]``, else None.""" + import numpy as np + + kind = _kind(values, micros) + nan = float("nan") + if kind == "i8": + return ( + kind, + np.array([0 if v is None else v for v in values], dtype="`` array of a column of float lists that float32 holds, or None.""" - import pyarrow as pa - - vec = _rows._float_vectors(values) - if vec is None: - return None - arr = pa.FixedSizeListArray.from_arrays(pa.array(vec.reshape(-1)), vec.shape[1]) - if any(v is None for v in values): - arr = pa.array(arr.to_pylist(), type=arr.type, mask=[v is None for v in values]) - return arr - - -def arrow_result(head_rows, columns, batches): +def arrow_result(head_rows, columns, batches, hints): """The ``pyarrow.Table`` of the buffered rows (dicts) followed by the decoded frames ``batches`` (a list of ``(count, chunks)``).""" - import pyarrow as pa - names = [] chunks_list = [] if head_rows: - table = _rows.to_arrow(head_rows, columns) - chunks = {} - for name in table.column_names: - column = table.column(name) - arr = ( - column.chunk(0) - if column.num_chunks == 1 - else pa.concat_arrays(column.chunks) - ) - if pa.types.is_null(arr.type): - # a frame types an all-null column as text - arr = pa.nulls(len(arr), pa.string()) - elif pa.types.is_list(arr.type) and arr.type.value_type == pa.float64(): - other = _first_part(batches, name) - if other is None or pa.types.is_fixed_size_list(other.type): - vector = _vector_arrow([row.get(name) for row in head_rows]) - if vector is not None: - arr = vector - chunks[name] = arr + frame, side = _head_frame(head_rows, columns, hints, False) + count, chunks, _ = _columnar.decode_arrow_frame(frame, side, _head_json_values) names.extend(chunks) - chunks_list.append((len(head_rows), chunks)) + chunks_list.append((count, chunks)) for count, chunks in batches: for name in chunks: if name not in names: names.append(name) chunks_list.append((count, chunks)) - return _columnar.merge_arrow(names, chunks_list) + return _columnar.merge_arrow(names, chunks_list, columns) diff --git a/src/arcadedb/_native/columns.py b/src/arcadedb/_native/columns.py deleted file mode 100644 index 90f59df6..00000000 --- a/src/arcadedb/_native/columns.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Rows to columns for the native backend (the columnar path is Python over the rows; the JPype path decodes a Java buffer).""" - -from datetime import date, datetime, timezone -from decimal import Decimal - - -def _names(rows, columns): - if columns is not None: - return list(columns) - seen = {} - for row in rows: - for k in row: - seen.setdefault(k, None) - return list(seen) - - -def _kind(values): - """'b', 'i', 'f', 'dt', 'dec', 'obj' for the non-null values of a column.""" - kinds = set() - for v in values: - if v is None: - continue - t = type(v) - if t is bool: - kinds.add("b") - elif t is int: - kinds.add("i") - elif t is float: - kinds.add("f") - elif t is datetime or t is date: - kinds.add("dt") - elif t is Decimal: - kinds.add("dec") - else: - kinds.add("obj") - if not kinds: - return "obj" - if len(kinds) == 1: - return kinds.pop() - if kinds == {"i", "f"}: - return "f" - return "obj" - - -def _ms(v): - if isinstance(v, datetime): - if v.tzinfo is not None: - v = v.astimezone(timezone.utc).replace(tzinfo=None) - return v - return datetime(v.year, v.month, v.day) - - -def _float_vectors(values): - """A column of equal-length float lists (None allowed) that float32 holds exactly, as a (count, dim) float32 array with - NaN rows for the nulls; None for any other column. Rows cannot say float[] from List, so only a column whose every - value survives float32 unchanged is read as a vector column (the JPype path decides on the Java element type). - """ - import numpy as np - - dim = 0 - for v in values: - if v is None: - continue - if type(v) is not list or not v or not all(type(x) is float for x in v): - return None - if dim == 0: - dim = len(v) - elif len(v) != dim: - return None - if dim == 0: - return None - arr = np.full((len(values), dim), np.nan, dtype=np.float32) - for i, v in enumerate(values): - if v is not None: - arr[i] = v - if not np.array_equal(arr[i].astype(np.float64), np.array(v)): - return None - return arr - - -def to_columns(rows, columns=None): - import numpy as np - - out = {} - for name in _names(rows, columns): - values = [row.get(name) for row in rows] - if _kind(values) == "obj": - vectors = _float_vectors(values) - if vectors is not None: - out[name] = vectors - continue - has_null = any(v is None for v in values) - kind = _kind(values) - if kind == "i": - out[name] = ( - np.array([np.nan if v is None else v for v in values], dtype=np.float64) - if has_null - else np.array(values, dtype=np.int64) - ) - elif kind == "f": - out[name] = np.array( - [np.nan if v is None else v for v in values], dtype=np.float64 - ) - elif kind == "b": - out[name] = values if has_null else np.array(values, dtype=bool) - elif kind == "dt": - out[name] = np.array( - [ - ( - np.datetime64("NaT", "ms") - if v is None - else np.datetime64(_ms(v), "ms") - ) - for v in values - ], - dtype="datetime64[ms]", - ) - elif kind == "dec": - arr = np.empty(len(values), dtype=object) - arr[:] = values - out[name] = arr - else: - out[name] = values - return out - - -def to_arrow(rows, columns=None): - import pyarrow as pa - - arrays = {} - for name in _names(rows, columns): - values = [row.get(name) for row in rows] - kind = _kind(values) - if kind == "i": - arrays[name] = pa.array(values, type=pa.int64()) - elif kind == "f": - try: - arrays[name] = pa.array(values, type=pa.float64()) - except ( - pa.ArrowInvalid - ): # an int beyond +-2**53 next to floats: the column degrades to exact text - arrays[name] = pa.array( - [None if v is None else str(v) for v in values], type=pa.string() - ) - elif kind == "b": - arrays[name] = pa.array(values, type=pa.bool_()) - elif kind == "dt": - arrays[name] = pa.array( - [None if v is None else _ms(v) for v in values], type=pa.timestamp("ms") - ) - elif kind == "dec": - try: - arrays[name] = pa.array(values) - except ( - pa.ArrowInvalid - ): # wider than decimal256 (76 digits): an exact string - arrays[name] = pa.array( - [None if v is None else str(v) for v in values], type=pa.string() - ) - else: - try: - arrays[name] = pa.array([None if v is None else v for v in values]) - except (pa.ArrowInvalid, pa.ArrowTypeError): - arrays[name] = pa.array( - [None if v is None else str(v) for v in values], type=pa.string() - ) - return pa.table(arrays) diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index 71d178eb..792b82ce 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -9,6 +9,7 @@ from datetime import date, datetime from typing import Any, Dict, Iterator, List, Optional +from .. import _columnar from .._params import refuse_surplus_positionals from ..exceptions import ArcadeDBError from ..transactions import TransactionContext @@ -210,9 +211,10 @@ def __init__(self, database, reply: dict, tagged: bool = True): self._rs = reply["rs"] self._buf = deque(zip(reply["rows"], reply["meta"])) self._more = not reply["drained"] - # the rows of this reply that the cursor still holds, for a columnar read to replay (arcadedb_n_next_col): the first reply of - # a cursor with more rows. Any later read of the cursor by row (arcadedb_n_next) lets the library drop them - self._head_n = len(reply["rows"]) if self._more else 0 + # the properties that were float[] ("f") or double[] ("d") in the rows read so far: a JSON row cannot tell them from a list of + # doubles, and a columnar read types them as vector columns (colframes.py) + self._vectors: Dict[str, str] = {} + self._note_vectors(reply) self._closed = False self._exhausted = False # the engine computes lazily: an error met while the first rows were read is raised when the rows before it @@ -227,7 +229,7 @@ def _lazy(cls, database, rs_id: int) -> "NativeResultSet": self._rs = rs_id self._buf = deque() self._tagged = False - self._head_n = 0 + self._vectors = {} self._more = rs_id > 0 self._closed = False self._exhausted = False @@ -262,7 +264,6 @@ def _fill(self) -> bool: raise _read_error(err) if not self._more: return False - self._head_n = 0 try: if lib.NEXT_BIN and _BINARY_REPLY: reply, tagged = loads_tagged( @@ -287,8 +288,15 @@ def _fill(self) -> bool: self._buf.extend(zip(reply["rows"], reply["meta"])) self._more = not reply["drained"] self._pending_error = reply.get("rerr") + self._note_vectors(reply) return bool(self._buf) or self._pending_error is not None + def _note_vectors(self, reply: dict) -> None: + for name in reply.get("fv", ()): + self._vectors[name] = "f" + for name in reply.get("dv", ()): + self._vectors[name] = "d" + def _finish(self) -> None: self._exhausted = True self.close() @@ -415,31 +423,24 @@ def iter_json_batches(self, batch_size: int = 10_000, **kwargs): def _read_frames(self, batch_size, columns, decode, flags=0): """The rows still to be read, in the two forms a result has: ``(head_rows, batches)``. - ``head_rows`` are rows that this object holds as dicts (the first reply, when the cursor no longer has them); - ``batches`` is the list of ``(count, parts)`` that ``decode(frame_bytes)`` makes of each column frame the cursor - yields (about 1 MiB of values each, ``batch_size`` rows at most; ``flags`` 1 asks for datetimes as microseconds). The result - is finished when this returns. + ``head_rows`` are the rows this object holds as dicts (the first reply, and any reply the caller read by row): they are + not in the cursor any more. ``batches`` is the list of ``(count, parts)`` that ``decode(frame_bytes)`` makes of each column + frame the cursor yields (about 1 MiB of values each, ``batch_size`` rows at most; ``flags`` 1 asks for datetimes as + microseconds). The result is finished when this returns. """ - if int(batch_size) < 1: - raise ValueError("batch_size must be at least 1") + size = _columnar.frame_rows(batch_size) head_rows: list = [] batches: list = [] if not self._readable(): return head_rows, batches spec = json.dumps(list(columns)).encode("utf-8") if columns is not None else b"" - skip = 0 - if self._head_n and self._more and self._rs: - # the cursor still holds the rows of the first reply: it replays the ones this object has not handed out - skip = self._head_n - len(self._buf) - self._buf.clear() - else: - head_rows = [props for props, _ in self._buf] - self._buf.clear() - if self._pending_error is not None: - self._fill() # the rows before the error are consumed: raise it + head_rows = [props for props, _ in self._buf] + self._buf.clear() + if self._pending_error is not None: + self._fill() # the rows before the error are consumed: raise it while self._more and self._rs: try: - raw = lib.call_frame(self._rs, int(batch_size), spec, flags, skip) + raw = lib.call_frame(self._rs, size, spec, flags) except lib.ColumnFrameError as exc: # the library closed the cursor self._more = False @@ -447,20 +448,29 @@ def _read_frames(self, batch_size, columns, decode, flags=0): self._closed = True self._exhausted = True raise _read_error(exc.reply) from None - skip = 0 - self._head_n = 0 count, parts, header = decode(raw) if header.get("done"): self._more = False # the library closed the cursor with the frame that ended the rows self._rs = 0 + elif count == 0: + # an open cursor must give a row per frame: never loop on a library that does not + self._more = False + self.close() + raise ArcadeDBError( + "The native library returned an empty column frame for a cursor with rows left (frame row limit " + f"{size}). This is a bug in the library or in the wheel." + ) if count: batches.append((count, parts)) self._finish() return head_rows, batches def to_dataframe(self, convert_types: bool = True): - """The rows as a ``pandas.DataFrame``: the values ``to_list()`` returns (a DATETIME keeps its microseconds, a DATE is a - ``datetime.date`` in an object column, a ``Decimal`` stays one), read as column frames, not as one dict per row. + """The rows as a ``pandas.DataFrame``, read as column frames (no dict per row). The values are the ones ``to_list()`` returns, + with two exceptions that are the same on both builds. A DATETIME keeps its microseconds (a ``datetime64[us]`` column, where + ``to_columns()`` and ``to_arrow()`` carry milliseconds), a DATE is a ``datetime.date`` in an object column, and a ``Decimal`` + stays one. A column of ``float[]`` vectors holds one numpy ``float32`` array per cell (a row of the matrix ``to_columns()`` + returns, not a Python list per row); ``to_list()`` gives lists of floats. """ import pandas as pd @@ -480,12 +490,17 @@ def _columns(self, batch_size, columns, micros): lambda raw: colframes.decode_columns(raw, micros), 1 if micros else 0, ) - return colframes.columns_result(head_rows, columns, micros, batches) + return colframes.columns_result( + head_rows, columns, micros, batches, self._vectors + ) def to_columns(self, batch_size: int = 25_000, columns=None): - """Columns as the JPype path returns them (numpy int64/float64/bool/datetime64[ms] arrays, lists for text and - JSON-typed values, an object array of Decimal), decoded from column frames by the decoder both builds share. None when - numpy is missing.""" + """Columns as the JVM build returns them (numpy int64/float64/bool/datetime64[ms] arrays, lists for text and JSON-typed + values, an object array of Decimal), decoded from column frames by the decoder both builds share. The arrays are writable, + and a column's type belongs to the whole result, not to where a batch ends. A column is a float32 matrix when the property is + a ``float[]`` (ARRAY_OF_FLOATS), and a list of lists when it is a LIST of numbers whatever the values. An empty result is + ``{}``, or the ``columns`` asked for as empty lists. None when numpy is missing. + """ try: __import__("numpy") # availability check except ImportError: @@ -505,7 +520,7 @@ def to_arrow(self, batch_size: int = 25_000, columns=None): head_rows, batches = self._read_frames( batch_size, columns, colframes.decode_arrow ) - return colframes.arrow_result(head_rows, columns, batches) + return colframes.arrow_result(head_rows, columns, batches, self._vectors) def to_polars(self, batch_size: int = 25_000, columns=None): try: diff --git a/src/arcadedb/_native/lib.py b/src/arcadedb/_native/lib.py index dc3f1e98..778305b5 100644 --- a/src/arcadedb/_native/lib.py +++ b/src/arcadedb/_native/lib.py @@ -232,7 +232,7 @@ class _Params(ctypes.Structure): "arcadedb_n_insert_colp": [P, L, S, S, P, I], "arcadedb_n_next": [P, L, I], "arcadedb_n_next_bin": [P, L, I], - "arcadedb_n_next_col": [P, L, I, S, I, I], + "arcadedb_n_next_col": [P, L, I, S, I], "arcadedb_n_rs_close": [P, L], "arcadedb_n_cmd": [P, L, I, S, S, S, S, L, I], "arcadedb_n_last_error": [P], @@ -366,15 +366,25 @@ def call_bin(name: str, *args): return raw.decode("utf-8"), floats -def call_frame(rs_id: int, max_rows: int, spec: bytes, flags: int, skip: int) -> bytes: +def call_frame(rs_id: int, max_rows: int, spec: bytes, flags: int) -> bytes: """One column frame of an open cursor (``arcadedb_n_next_col``): its bytes, in the layout ``arcadedb._columnar`` reads. The reply is an int64 byte count and the frame. A negative count is the negated length of an error JSON: it is raised as - :class:`ColumnFrameError` with the decoded reply (the cursor is already closed on the Java side then). + :class:`ColumnFrameError` with the decoded reply (the cursor is already closed on the Java side then). ``max_rows`` is a C int: + callers clamp it (``arcadedb._columnar.frame_rows``). """ lib = _lib or _load() + if getattr(lib, "arcadedb_n_next_col", None) is None: + raise ArcadeDBError( + "The native library is older than this wheel: it has no arcadedb_n_next_col entry point, which to_arrow, to_columns, " + "to_dataframe and to_polars need. Reinstall the arcadedb wheel (a library and its Python files ship together)." + ) th = _thread_handle() - ptr = lib.arcadedb_n_next_col(th, rs_id, max_rows, spec, flags, skip) + ptr = lib.arcadedb_n_next_col(th, rs_id, max_rows, spec, flags) + if not ptr: + raise ArcadeDBError( + "The native library returned no reply to a column frame request" + ) try: n = int.from_bytes(ctypes.string_at(ptr, 8), "little", signed=True) if n < 0: diff --git a/src/arcadedb/results.py b/src/arcadedb/results.py index 945ab21e..2abc95dc 100644 --- a/src/arcadedb/results.py +++ b/src/arcadedb/results.py @@ -15,6 +15,9 @@ from ._columnar import ( decode_arrow_frame, decode_columns_frame, + empty_arrow, + empty_columns, + frame_rows, merge_arrow, merge_columns, ) @@ -546,10 +549,14 @@ def to_dataframe(self, convert_types: bool = True): Requires pandas to be installed. - Values are the ones ``to_list()`` returns: a DATETIME keeps its - microseconds (a ``datetime64[us]`` column, where ``to_columns()`` and - ``to_arrow()`` carry milliseconds), a DATE is a ``datetime.date`` in an - object column, and a ``Decimal`` stays one. + Values are the ones ``to_list()`` returns, with two exceptions that are + the same on both builds. A DATETIME keeps its microseconds (a + ``datetime64[us]`` column, where ``to_columns()`` and ``to_arrow()`` + carry milliseconds), a DATE is a ``datetime.date`` in an object column, + and a ``Decimal`` stays one. A column of ``float[]`` vectors holds one + numpy ``float32`` array per cell (a row of the matrix ``to_columns()`` + returns, no copy into a Python list per row); ``to_list()`` gives + lists of floats. Args: convert_types: Convert Java types to Python (default: True) @@ -599,7 +606,9 @@ def to_polars( Reads :meth:`to_arrow` and hands the table to Polars, so nulls keep their type (a nullable integer stays an integer column) and a DECIMAL - column is a Polars ``Decimal``. Requires polars and pyarrow + column is a Polars ``Decimal``. A column of ``float[]`` vectors is a + Polars ``Array(Float32, n)`` (``to_arrow()`` types it as a fixed-size + list of float32). Requires polars and pyarrow (``uv add 'arcadedb-jvm[polars]'``). Args: @@ -645,7 +654,11 @@ def to_columns( Null handling follows pandas conventions: int/datetime columns with nulls are promoted to float64 with NaN / datetime64 NaT; string and JSON columns use None. A DECIMAL column is an object array of exact - ``Decimal`` values (#115). + ``Decimal`` values (#115). The type of a column belongs to the whole + result, not to where a batch ends: a batch that lacks the column or + holds only nulls for it takes the type of the batches that have + values. The arrays are writable. An empty result is ``{}``, or, when + ``columns`` is given, those columns as empty lists. The columns are the union of every row's property names, in order of first appearance, because a document is schemaless: a property the @@ -690,13 +703,13 @@ def _read_columns(self, batch_size, columns, micros): # ";".join: a name may legally contain a semicolon, a quote or a # backslash. spec = json.dumps(list(columns)) if columns is not None else "" + size = frame_rows(batch_size) if not self._readable(): - return {} - total = 0 + return empty_columns(columns) while True: try: java_batch = column_batcher.nextColumnBatch( - self._java_result_set, int(batch_size), spec, micros + self._java_result_set, size, spec, micros ) except JException as exc: raise _read_error(exc) from exc @@ -711,11 +724,8 @@ def _read_columns(self, batch_size, columns, micros): if name not in names: names.append(name) batches.append((count, parts)) - total += count - if total == 0: - return {} - return merge_columns(names, batches) + return merge_columns(names, batches, columns) def to_arrow( self, batch_size: int = 25_000, columns: Optional[Sequence[str]] = None @@ -748,7 +758,7 @@ def to_arrow( """ try: __import__("numpy") # availability check - import pyarrow as pa + __import__("pyarrow") # availability check except ImportError: return None @@ -763,13 +773,13 @@ def to_arrow( batches: List[Tuple[int, Dict[str, Any]]] = [] spec = json.dumps(list(columns)) if columns is not None else "" + size = frame_rows(batch_size) if not self._readable(): - return pa.table({}) - total = 0 + return empty_arrow(columns) while True: try: java_batch = column_batcher.nextColumnBatch( - self._java_result_set, int(batch_size), spec, False + self._java_result_set, size, spec, False ) except JException as exc: raise _read_error(exc) from exc @@ -784,11 +794,8 @@ def to_arrow( if name not in names: names.append(name) batches.append((count, chunks)) - total += count - if total == 0 or not names: - return pa.table({}) - return merge_arrow(names, batches) + return merge_arrow(names, batches, columns) def iter_chunks( self, size: int = 1000, convert_types: bool = True diff --git a/src/java/com/arcadedb/python/ColumnBatcher.java b/src/java/com/arcadedb/python/ColumnBatcher.java index cf971a20..168b2851 100644 --- a/src/java/com/arcadedb/python/ColumnBatcher.java +++ b/src/java/com/arcadedb/python/ColumnBatcher.java @@ -151,16 +151,22 @@ public static Object[] nextColumnBatch(final ResultSet rs, final int max, final final byte[] nulls = new byte[(count + 7) / 8]; // detect the column type across ALL values; mixed/unsupported -> json String type = null; + boolean sawInt = false, sawFloating = false, bigInt = false; for (int r = 0; r < count; r++) { final Object v = rows.get(r)[c]; if (v == null) continue; final String t; - if (v instanceof Long || v instanceof Integer || v instanceof Short || v instanceof Byte) + if (v instanceof Long || v instanceof Integer || v instanceof Short || v instanceof Byte) { t = "i8"; - else if (v instanceof Double || v instanceof Float) + sawInt = true; + final long l = ((Number) v).longValue(); + if (l > (1L << 53) || l < -(1L << 53)) + bigInt = true; + } else if (v instanceof Double || v instanceof Float) { t = "f8"; - else if (v instanceof Boolean) + sawFloating = true; + } else if (v instanceof Boolean) t = "b1"; else if (v instanceof Date || v instanceof Calendar || v instanceof LocalDateTime || v instanceof LocalDate || v instanceof Instant @@ -198,6 +204,9 @@ else if (!type.equals(t)) { } if (type == null) type = "str"; // all-null column + // an integer beyond 2**53 next to a floating value has no exact double: the values stay what they are, in a json column + if (type.equals("f8") && sawInt && sawFloating && bigInt) + type = "json"; int vdim = -1; if (type.equals("f4v") || type.equals("f8v")) { diff --git a/tests/test_column_frames.py b/tests/test_column_frames.py index 7956b15a..dc06f830 100644 --- a/tests/test_column_frames.py +++ b/tests/test_column_frames.py @@ -336,8 +336,198 @@ def test_to_dataframe_and_to_polars(table): ] assert [list(v) for v in df["v3"]] == [r["v3"] for r in rows] assert df["ts"].iloc[1] == datetime(2026, 1, 2, 3, 4, 5, 1000) + assert str(df["ts"].dtype).startswith( + "datetime64[us" + ) # microseconds on both wheels, whatever pandas infers from datetimes + # a vector cell is a numpy float32 array, not a Python list (documented on to_dataframe) + assert all(isinstance(v, np.ndarray) and v.dtype == np.float32 for v in df["v3"]) import polars as pl p = db.query("sql", _sel()).to_polars() assert isinstance(p, pl.DataFrame) and p.height == 5000 + assert p.schema["v3"] == pl.Array(pl.Float32, 3) assert p["big"].to_list() == [r["big"] for r in rows] + + +# --- what a column's type depends on: the whole result, not the size of the result or where a batch ends --------------------- + + +@pytest.mark.parametrize("batch_size", [2**31 - 1, 2**31, 2**32, 10**10]) +def test_a_huge_batch_size_means_as_many_as_one_batch_takes(table, batch_size): + db, rows = table + t = db.query("sql", _sel()).to_arrow(batch_size=batch_size) + assert t.num_rows == 5000 + c = db.query("sql", _sel()).to_columns(batch_size=batch_size) + assert len(c["big"]) == 5000 + + +@pytest.mark.parametrize("batch_size", [0, -1]) +def test_a_batch_size_below_one_is_refused(table, batch_size): + db, _ = table + with pytest.raises(ValueError): + db.query("sql", _sel()).to_arrow(batch_size=batch_size) + with pytest.raises(ValueError): + db.query("sql", _sel()).to_columns(batch_size=batch_size) + + +def test_no_rows_with_columns_given_returns_those_columns_empty(table): + db, _ = table + sql = "SELECT n, s FROM T WHERE big < 0" + t = db.query("sql", sql).to_arrow(columns=["s", "n"]) + assert t.column_names == ["s", "n"] and t.num_rows == 0 + assert db.query("sql", sql).to_columns(columns=["s", "n"]) == {"s": [], "n": []} + assert db.query("sql", sql).to_columns() == {} + assert db.query("sql", sql).to_arrow().num_columns == 0 + + +@pytest.mark.parametrize("limit", [1, 10, 300]) +def test_to_columns_arrays_are_writable(table, limit): + db, _ = table + c = db.query( + "sql", + f"SELECT big, d, ts, v3 FROM T ORDER BY big LIMIT {limit}", # nosec B608 - an integer of the test + ).to_columns() + for name in ("big", "d", "ts", "v3"): + assert c[name].flags.writeable, (name, limit) + c[name][0] = c[name][0] + + +@pytest.mark.parametrize("rows", [255, 256, 257, 511, 512, 513]) +def test_sizes_around_the_first_reply_and_two_of_them(table, rows): + db, stored = table + t = db.query("sql", _sel() + f" LIMIT {rows}").to_arrow() + assert t.num_rows == rows + assert t.column("big").to_pylist() == [r["big"] for r in stored[:rows]] + assert t.column("v9").to_pylist() == [r["v9"] for r in stored[:rows]] + assert t.schema.field("v9").type == pa.list_(pa.float32(), 9) + assert t.schema.field("s").type == pa.string() + + +def test_two_cursors_open_at_once_and_a_cursor_closed_early(table): + db, rows = table + one = db.query("sql", _sel()) + two = db.query("sql", _sel()) + assert next(one).get("big") == rows[0]["big"] + assert next(two).get("big") == rows[0]["big"] + t_two = two.to_arrow() + t_one = one.to_arrow() + assert t_one.column("big").to_pylist() == [r["big"] for r in rows[1:]] + assert t_two.column("big").to_pylist() == [r["big"] for r in rows[1:]] + partly = db.query("sql", _sel()) + for _ in range(300): + next(partly) + partly.close() + fresh = db.query("sql", _sel() + " LIMIT 400").to_arrow() + assert fresh.num_rows == 400 + + +def test_a_cursor_with_nothing_read_closed_many_times_over(table): + db, _ = table + cursors = [db.query("sql", _sel()) for _ in range(300)] + for cursor in cursors: + cursor.close() + assert db.query("sql", _sel() + " LIMIT 3").to_arrow().num_rows == 3 + + +@pytest.fixture +def declared(temp_db): + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE Dec") + db.command("sql", "CREATE PROPERTY Dec.i INTEGER") + db.command("sql", "CREATE PROPERTY Dec.arr ARRAY_OF_FLOATS") + db.command("sql", "CREATE PROPERTY Dec.dbl ARRAY_OF_DOUBLES") + db.command("sql", "CREATE PROPERTY Dec.lst LIST") + rows = [ + { + "i": i, + "arr": [0.5 * i, 1.0, 2.0], + "dbl": [0.5 * i, 1.5, 2.5], + "lst": [0.5 * i, 1.0, 2.0], + } + for i in range(3000) + ] + db.insert_many("Dec", rows) + return db, rows + + +@pytest.mark.parametrize("rows", [10, 300, 3000]) +def test_a_vector_column_is_a_declared_array_whatever_the_result_size(declared, rows): + db, stored = declared + sql = f"SELECT i, arr, dbl, lst FROM Dec ORDER BY i LIMIT {rows}" # nosec B608 - an integer of the test + c = db.query("sql", sql).to_columns() + assert isinstance(c["arr"], np.ndarray) and c["arr"].dtype == np.float32 + assert isinstance(c["dbl"], np.ndarray) and c["dbl"].dtype == np.float64 + assert isinstance(c["lst"], list), type( + c["lst"] + ) # a LIST of doubles is lists, not a float32 matrix that happens to fit + assert c["lst"] == [r["lst"] for r in stored[:rows]] + t = db.query("sql", sql).to_arrow() + assert t.schema.field("arr").type == pa.list_(pa.float32(), 3) + assert t.schema.field("dbl").type == pa.list_(pa.float64(), 3) + assert t.schema.field("lst").type == pa.list_(pa.float64()) + assert t.column("dbl").to_pylist() == [r["dbl"] for r in stored[:rows]] + + +def test_a_double_array_stays_double_after_rows_read_past_the_first_reply(declared): + db, stored = declared + rs = db.query("sql", "SELECT i, dbl, arr FROM Dec ORDER BY i") + for _ in range(300): + next(rs) + t = rs.to_arrow() + assert t.schema.field("dbl").type == pa.list_(pa.float64(), 3) + assert t.schema.field("arr").type == pa.list_(pa.float32(), 3) + assert t.column("dbl").to_pylist() == [r["dbl"] for r in stored[300:]] + + +def test_a_property_that_is_numeric_from_row_30000_on_has_one_type(temp_db): + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE Late") + db.command("sql", "CREATE PROPERTY Late.i INTEGER") + rows = [{"i": i} for i in range(30000)] + [ + {"i": i, "late": i * 2} for i in range(30000, 30100) + ] + db.insert_many("Late", rows) + sql = "SELECT i, late FROM Late ORDER BY i" + t = db.query("sql", sql).to_arrow(batch_size=25_000, columns=["i", "late"]) + assert t.schema.field("late").type == pa.int64() + late = t.column("late").to_pylist() + assert late[:30000] == [None] * 30000 and late[30000:] == [ + i * 2 for i in range(30000, 30100) + ] + c = db.query("sql", sql).to_columns(batch_size=25_000, columns=["i", "late"]) + assert isinstance(c["late"], np.ndarray) and c["late"].dtype == np.float64 + assert np.isnan(c["late"][:30000]).all() + assert list(c["late"][30000:]) == [i * 2 for i in range(30000, 30100)] + + +@pytest.mark.parametrize("rows", [10, 300]) +def test_an_integer_beyond_2_53_next_to_a_float_stays_exact(temp_db, rows): + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE Mix") + db.command("sql", "CREATE PROPERTY Mix.i INTEGER") + big = 2**60 + 1 + db.insert_many( + "Mix", [{"i": i, "x": big if i % 2 == 0 else 1.5} for i in range(rows)] + ) + c = db.query("sql", "SELECT i, x FROM Mix ORDER BY i").to_columns() + want = [big if i % 2 == 0 else 1.5 for i in range(rows)] + assert list(c["x"]) == want + assert all(type(a) is type(b) for a, b in zip(c["x"], want)) + t = db.query("sql", "SELECT i, x FROM Mix ORDER BY i").to_arrow() + assert t.column("x").to_pylist() == [ + str(v) for v in want + ] # no common numeric type: exact text, as across batches + + +def test_vectors_whose_dimension_changes_between_frames_are_one_list_column(temp_db): + db = temp_db + db.command("sql", "CREATE DOCUMENT TYPE Dim") + db.command("sql", "CREATE PROPERTY Dim.i INTEGER") + db.command("sql", "CREATE PROPERTY Dim.v ARRAY_OF_FLOATS") + db.insert_many( + "Dim", + [{"i": i, "v": [0.5, 1.0, 2.0] if i < 300 else [0.5, 1.0]} for i in range(600)], + ) + t = db.query("sql", "SELECT i, v FROM Dim ORDER BY i").to_arrow(batch_size=100) + assert t.schema.field("v").type == pa.list_(pa.float32()) + assert t.column("v").to_pylist() == [[0.5, 1.0, 2.0]] * 300 + [[0.5, 1.0]] * 300 diff --git a/tests/test_native_column_frames_reference.py b/tests/test_native_column_frames_reference.py index 630a4390..2fc38c32 100644 --- a/tests/test_native_column_frames_reference.py +++ b/tests/test_native_column_frames_reference.py @@ -262,23 +262,20 @@ def ref_db(tmp_path_factory): db.close() -def _norm(v): - """A value as a plain Python object, NaN, NaT and None alike (pandas turns a None into NaN in some columns).""" - if isinstance(v, np.ndarray): - return [_norm(x) for x in v.tolist()] - if isinstance(v, (list, tuple)): - return [_norm(x) for x in v] - if v is None or v is pd.NaT: - return None - if isinstance(v, float) and math.isnan(v): - return None - if isinstance(v, np.datetime64): - return None if np.isnat(v) else v.astype("datetime64[us]").item() - if isinstance(v, np.generic): - return v.item() - if isinstance(v, pd.Timestamp): - return v.to_pydatetime() - return v +def _eq(have, want): + """Equal values, NaN equal to NaN, None equal only to None: a missing value is not a NaN.""" + if isinstance(have, np.ndarray) or isinstance(want, np.ndarray): + return np.array_equal(np.asarray(have), np.asarray(want), equal_nan=True) + if isinstance(have, (list, tuple)) and isinstance(want, (list, tuple)): + return len(have) == len(want) and all(_eq(x, y) for x, y in zip(have, want)) + if ( + isinstance(have, float) + and isinstance(want, float) + and math.isnan(have) + and math.isnan(want) + ): + return True + return have == want def _old_rows(db, sql, *args): @@ -292,16 +289,19 @@ def _assert_arrow_equal(new, old): have, want = new.column(name), old.column(name) if pa.types.is_null(want.type): assert have.type == pa.string() and have.null_count == len(have), name - elif pa.types.is_list(want.type) and want.type.value_type == pa.float64(): - assert pa.types.is_fixed_size_list(have.type) or pa.types.is_list( - have.type + elif ( + pa.types.is_list(want.type) + and want.type.value_type == pa.float64() + and name.startswith("v") + ): + # a float[] column: the same values as a fixed-size list of float32 (a float32 widened to float64 is exact) + assert have.type == pa.list_( + pa.float32(), + len(want[0].as_py() if want[0].is_valid else want[1].as_py()), ), name - assert have.type.value_type in (pa.float32(), pa.float64()), name else: assert have.type == want.type, (name, have.type, want.type) - assert [_norm(x) for x in have.to_pylist()] == [ - _norm(x) for x in want.to_pylist() - ], name + assert _eq(have.to_pylist(), want.to_pylist()), name def _assert_columns_equal(new, old): @@ -312,9 +312,22 @@ def _assert_columns_equal(new, old): assert isinstance(have, np.ndarray), (name, type(have)) assert have.dtype == want.dtype, (name, have.dtype, want.dtype) assert have.shape == want.shape, name + np.testing.assert_array_equal(have, want, err_msg=name) else: assert isinstance(have, list), (name, type(have)) - assert _norm(have) == _norm(want), name + assert _eq(have, want), name + + +def _missing_as_none(values): + """pandas decides how a missing value is shown (None, nan or NaT, by dtype and version): all of them are one here.""" + return [ + ( + None + if v is None or v is pd.NaT or (isinstance(v, float) and math.isnan(v)) + else v + ) + for v in values + ] def _assert_frame_equal(new, old_rows, columns=None): @@ -323,7 +336,31 @@ def _assert_frame_equal(new, old_rows, columns=None): old = old.reindex(columns=list(columns)) assert list(new.columns) == list(old.columns) for name in old.columns: - assert _norm(new[name].tolist()) == _norm(old[name].tolist()), name + if name.startswith("v") and name != "vv": + # a vector cell is a numpy float32 array (documented), the values are the stored ones + for cell, want in zip(new[name].tolist(), old[name].tolist()): + if want is None: + assert cell is None or ( + isinstance(cell, float) and math.isnan(cell) + ), name + else: + assert isinstance(cell, np.ndarray) and cell.dtype == np.float32, ( + name, + type(cell), + ) + assert _eq(cell.tolist(), want), name + continue + have = _missing_as_none(new[name].tolist()) + want = _missing_as_none(old[name].tolist()) + assert [ + x.to_pydatetime() if isinstance(x, pd.Timestamp) else x for x in have + ] == [ + x.to_pydatetime() if isinstance(x, pd.Timestamp) else x for x in want + ], name + if "ts" in old.columns: + assert str(new["ts"].dtype).startswith("datetime64[us"), new[ + "ts" + ].dtype # microseconds, whatever pandas infers from datetimes QUERIES = [ @@ -425,10 +462,68 @@ def test_to_list_still_reads_the_json_rows(ref_db): assert len(rows) == 300 and rows[299]["big"] == (1 << 40) + 299 -def test_the_cursor_keeps_no_rows_after_a_row_read(ref_db): +def test_a_frame_stays_near_its_budget_when_every_row_has_other_properties(tmp_path): + """Rows that each carry a property no other row has: the frame holds rows times the union of the names, not the values read.""" + from unittest import mock + + import arcadedb + from arcadedb._native import lib + + db = arcadedb.create_database(str(tmp_path / "sparse")) + try: + db.command("sql", "CREATE DOCUMENT TYPE Sparse") + db.insert_many("Sparse", [{"i": i, f"p{i}": i} for i in range(2000)]) + sizes = [] + real = lib.call_frame + + def spy(*args): + raw = real(*args) + sizes.append(len(raw)) + return raw + + with mock.patch.object(lib, "call_frame", spy): + t = db.query("sql", "SELECT FROM Sparse ORDER BY i").to_arrow() + assert t.num_rows == 2000 and t.num_columns == 2001 + assert t.column("p1999").to_pylist() == [None] * 1999 + [1999] + # about 1 MiB of fixed-width buffers (rows x columns x 9 bytes) is what a frame is allowed; 2000 x 2000 x 9 = 36 MiB in one frame is not + assert len(sizes) > 5 and max(sizes) < 3 * (1 << 20), sizes + finally: + db.close() + + +def test_an_empty_frame_for_a_cursor_with_rows_left_is_an_error_not_a_loop(ref_db): + from unittest import mock + + from arcadedb._native import lib + from arcadedb.exceptions import ArcadeDBError + + header = b'{"count":0,"done":false,"cols":[]}' + empty = len(header).to_bytes(4, "little") + header + rs = ref_db.query("sql", "SELECT FROM R ORDER BY big") - assert rs._head_n == 256 - for _ in range(300): - next(rs) - assert rs._head_n == 0 - rs.close() + with mock.patch.object(lib, "call_frame", lambda *args: empty): + with pytest.raises(ArcadeDBError, match="empty column frame"): + rs.to_arrow() + + +def test_a_cursor_reads_a_huge_batch_size_without_wrapping(ref_db): + sql = "SELECT big FROM R ORDER BY big" + for batch_size in (2**31, 2**32, 2**33 + 5): + assert ref_db.query("sql", sql).to_arrow(batch_size=batch_size).num_rows == ROWS + + +def test_a_list_property_is_not_a_vector_column_at_any_result_size(tmp_path): + import arcadedb + + db = arcadedb.create_database(str(tmp_path / "lists")) + try: + db.command("sql", "CREATE DOCUMENT TYPE L") + db.insert_many("L", [{"i": i, "lst": [0.5, 1.0, 2.0]} for i in range(600)]) + for rows in (5, 300, 600): + c = db.query( + "sql", + f"SELECT i, lst FROM L ORDER BY i LIMIT {rows}", # nosec B608 - an integer of the test + ).to_columns() + assert isinstance(c["lst"], list) and c["lst"] == [[0.5, 1.0, 2.0]] * rows + finally: + db.close() From 17ccea5b3ba3050e7a1cbc92546962bd8aba2db0 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Mon, 9 Nov 2026 22:23:03 +0900 Subject: [PATCH 5/6] Native frames: vector types unify to float64, a null vector is None in to_dataframe A dimension that changes inside a frame makes a json column (list), so the same change between two frames now gives list too (it gave list); float32 next to float64 of one size is fixed_size_list. In to_dataframe a vector cell is a numpy float32 array or None for a null, on both wheels (a NaN row, which to_columns keeps in its matrix, read as a vector of NaN where the old route had a missing value); to_arrow and to_polars already kept the null. Co-Authored-By: Claude Sonnet 5.5 --- src/arcadedb/_columnar.py | 16 +++-- src/arcadedb/_native/database.py | 9 +-- src/arcadedb/results.py | 15 ++--- tests/test_column_frames.py | 65 +++++++++++++++++++- tests/test_native_column_frames_reference.py | 5 +- 5 files changed, 87 insertions(+), 23 deletions(-) diff --git a/src/arcadedb/_columnar.py b/src/arcadedb/_columnar.py index 16c66d81..75dd8f5e 100644 --- a/src/arcadedb/_columnar.py +++ b/src/arcadedb/_columnar.py @@ -170,14 +170,15 @@ def _unify_arrow_chunk_types(arrs: list, pa) -> list: # (the widening is exact), never text values = {t.value_type for t in types} if values <= {pa.float32(), pa.float64()}: - value = pa.float64() if pa.float64() in values else pa.float32() + # always float64: the widening is exact, and a frame whose rows differ in length is a json column (list) on its + # own, so float32 here would make the type depend on whether the change falls inside a frame or between two sizes = { t.list_size if pa.types.is_fixed_size_list(t) else None for t in types } target = ( - pa.list_(value, next(iter(sizes))) + pa.list_(pa.float64(), next(iter(sizes))) if len(sizes) == 1 and None not in sizes - else pa.list_(value) + else pa.list_(pa.float64()) ) try: return [pa.array(a.to_pylist(), type=target) for a in arrs] @@ -238,7 +239,10 @@ def decode_columns_frame( ) -> Tuple[int, Dict[str, Any], Dict[str, Any]]: """``(count, parts, header)`` of one frame as ``to_columns`` columns: ``parts`` maps a column name to a numpy array (int64, float64, bool, datetime64) or a Python list (text, JSON-typed values, a bool column with nulls); a count of 0 has - no parts. With ``micros`` a naive datetime column is ``datetime64[us]`` (``dtu``). + no parts. A vector column is a (count, dim) matrix, a null row of it NaN. + + ``micros`` is the ``to_dataframe`` form: a naive datetime column is ``datetime64[us]`` (``dtu``), and a vector column is a + list with one numpy array per cell and None for a null (a NaN row would read as a vector of NaN, not as a missing value). """ import numpy as np @@ -291,6 +295,10 @@ def decode_columns_frame( values = np.empty((count, 0), dtype=dt) else: values = np.frombuffer(data, dtype=dt).reshape(count, dim).copy() + if micros: + values = [ + None if has_nulls and mask[i] else values[i] for i in range(count) + ] elif ctype == "json": values = json_values(data, side, count) elif ctype == "dec": diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index 792b82ce..4cd943d3 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -470,15 +470,11 @@ def to_dataframe(self, convert_types: bool = True): with two exceptions that are the same on both builds. A DATETIME keeps its microseconds (a ``datetime64[us]`` column, where ``to_columns()`` and ``to_arrow()`` carry milliseconds), a DATE is a ``datetime.date`` in an object column, and a ``Decimal`` stays one. A column of ``float[]`` vectors holds one numpy ``float32`` array per cell (a row of the matrix ``to_columns()`` - returns, not a Python list per row); ``to_list()`` gives lists of floats. + returns, not a Python list per row) and None for a null; ``to_list()`` gives lists of floats. """ import pandas as pd columns = self._columns(25_000, None, True) - # pandas rejects 2-D ndarrays as column values; vector columns become object columns of row slices - columns = { - k: (list(v) if getattr(v, "ndim", 1) > 1 else v) for k, v in columns.items() - } return pd.DataFrame(columns) def _columns(self, batch_size, columns, micros): @@ -498,7 +494,8 @@ def to_columns(self, batch_size: int = 25_000, columns=None): """Columns as the JVM build returns them (numpy int64/float64/bool/datetime64[ms] arrays, lists for text and JSON-typed values, an object array of Decimal), decoded from column frames by the decoder both builds share. The arrays are writable, and a column's type belongs to the whole result, not to where a batch ends. A column is a float32 matrix when the property is - a ``float[]`` (ARRAY_OF_FLOATS), and a list of lists when it is a LIST of numbers whatever the values. An empty result is + a ``float[]`` (ARRAY_OF_FLOATS), and a list of lists when it is a LIST of numbers whatever the values; a null vector is a row of + NaN in the matrix (``to_dataframe()`` and ``to_arrow()`` keep it a null). An empty result is ``{}``, or the ``columns`` asked for as empty lists. None when numpy is missing. """ try: diff --git a/src/arcadedb/results.py b/src/arcadedb/results.py index 2abc95dc..369c1dfa 100644 --- a/src/arcadedb/results.py +++ b/src/arcadedb/results.py @@ -555,8 +555,8 @@ def to_dataframe(self, convert_types: bool = True): carry milliseconds), a DATE is a ``datetime.date`` in an object column, and a ``Decimal`` stays one. A column of ``float[]`` vectors holds one numpy ``float32`` array per cell (a row of the matrix ``to_columns()`` - returns, no copy into a Python list per row); ``to_list()`` gives - lists of floats. + returns, no copy into a Python list per row) and None for a null; + ``to_list()`` gives lists of floats. Args: convert_types: Convert Java types to Python (default: True) @@ -588,12 +588,7 @@ def to_dataframe(self, convert_types: bool = True): # to_arrow() keep the documented milliseconds. columns = self._read_columns(25_000, None, micros=True) if columns is not None: - # pandas rejects 2-D ndarrays as column values; vector - # columns (f4v/f8v) become object columns of row slices - columns = { - k: (list(v) if getattr(v, "ndim", 1) > 1 else v) - for k, v in columns.items() - } + # a vector column is already one numpy array per cell, None for a null return pd.DataFrame(columns) return pd.DataFrame(self.to_list(convert_types=convert_types)) @@ -657,7 +652,9 @@ def to_columns( ``Decimal`` values (#115). The type of a column belongs to the whole result, not to where a batch ends: a batch that lacks the column or holds only nulls for it takes the type of the batches that have - values. The arrays are writable. An empty result is ``{}``, or, when + values. The arrays are writable. A vector column is a ``(count, dim)`` matrix + and a null vector is a row of NaN in it (``to_dataframe()`` and ``to_arrow()`` keep + it a null). An empty result is ``{}``, or, when ``columns`` is given, those columns as empty lists. The columns are the union of every row's property names, in order of diff --git a/tests/test_column_frames.py b/tests/test_column_frames.py index dc06f830..177c2d79 100644 --- a/tests/test_column_frames.py +++ b/tests/test_column_frames.py @@ -529,5 +529,68 @@ def test_vectors_whose_dimension_changes_between_frames_are_one_list_column(temp [{"i": i, "v": [0.5, 1.0, 2.0] if i < 300 else [0.5, 1.0]} for i in range(600)], ) t = db.query("sql", "SELECT i, v FROM Dim ORDER BY i").to_arrow(batch_size=100) - assert t.schema.field("v").type == pa.list_(pa.float32()) + assert t.schema.field("v").type == pa.list_( + pa.float64() + ) # a lossless widening, wherever the dimension changes assert t.column("v").to_pylist() == [[0.5, 1.0, 2.0]] * 300 + [[0.5, 1.0]] * 300 + + +def test_vector_types_unify_by_value_and_dimension_never_to_text(): + def vec(dim, ctype, dtype, rows=2): + data = np.arange(rows * dim, dtype=dtype).reshape(rows, dim).tobytes() + return _frame(rows, [("v", ctype, [], data, dim)]) + + f3, f2 = _arrow(vec(3, "f4v", " Date: Mon, 9 Nov 2026 23:01:07 +0900 Subject: [PATCH 6/6] Native frames: a null vector is a NaN row in the to_columns matrix whatever the writer put there The decoder set no value for the null rows of a vector column and read what the frame held, which is zeros in a hand-built frame and NaN only where the writer chose NaN; a zero row reads as a real zero vector. decode_columns_frame now sets the null rows to NaN itself, so the single-frame matrix, the matrix after a merge with other frames and the first-reply frame agree, and the array stays writable. Co-Authored-By: Claude Sonnet 5.5 --- src/arcadedb/_columnar.py | 6 ++++-- tests/test_column_frames.py | 13 ++++++++++++- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/arcadedb/_columnar.py b/src/arcadedb/_columnar.py index 75dd8f5e..c33ff0e8 100644 --- a/src/arcadedb/_columnar.py +++ b/src/arcadedb/_columnar.py @@ -287,14 +287,16 @@ def decode_columns_frame( else: values = arr elif ctype in ("f4v", "f8v"): - # fixed-dimension vector column -> 2-D array (count, dim); the writer puts NaN rows into null slots; - # an empty vector has no "dim" in the header + # fixed-dimension vector column -> 2-D array (count, dim); an empty vector has no "dim" in the header. The null rows are + # made NaN here, whatever the writer put in the slot: a row of zeros would read as a real zero vector dim = col.get("dim", 0) dt = "