diff --git a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java index 9b5acb91..13ce33a2 100644 --- a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java +++ b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java @@ -317,6 +317,52 @@ 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). 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) { + ResultSet rs = null; + try { + rs = RESULT_SETS.get(rsId); + if (rs == null) + throw new IllegalArgumentException("unknown result set " + 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, max, pinned, (flags & 1) != 0); + if (frame.done) { + RESULT_SETS.remove(rsId); + closeQuietly(rs); + } + return frameOut(frame.bytes.length, frame.bytes); + } catch (final Throwable t) { + RESULT_SETS.remove(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)); 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..02d9106c --- /dev/null +++ b/native/src/main/java/org/humemai/nativelib/ColumnFrames.java @@ -0,0 +1,380 @@ +/* + * 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 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 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() < 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 (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; + 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]); + } + } + 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); + 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..b3bbab06 100644 --- a/native/src/main/java/org/humemai/nativelib/NativeCodec.java +++ b/native/src/main/java/org/humemai/nativelib/NativeCodec.java @@ -51,8 +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). */ + /** + * 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(']'); + } + } + 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; @@ -71,7 +101,7 @@ 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); @@ -90,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()) { @@ -102,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 new file mode 100644 index 00000000..c33ff0e8 --- /dev/null +++ b/src/arcadedb/_columnar.py @@ -0,0 +1,471 @@ +"""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_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()}: + # 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_(pa.float64(), next(iter(sizes))) + if len(sizes) == 1 and None not in sizes + else pa.list_(pa.float64()) + ) + 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: + return [a.cast(common) for a in arrs] + + 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 + + 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. 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 + + 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": + # a copy: np.frombuffer over bytes is read-only, and every array of to_columns is writable + arr = np.frombuffer(data, dtype=" 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 = " 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]]], columns=None +) -> Dict[str, Any]: + """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: + 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 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)] + 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=" 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=", 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 f30dc526..4cd943d3 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,6 +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 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 @@ -224,6 +229,7 @@ def _lazy(cls, database, rs_id: int) -> "NativeResultSet": self._rs = rs_id self._buf = deque() self._tagged = False + self._vectors = {} self._more = rs_id > 0 self._closed = False self._exhausted = False @@ -282,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() @@ -407,33 +420,104 @@ 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 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. + """ + 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"" + 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, size, spec, flags) + 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 + 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``, 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) and None for a null; ``to_list()`` gives lists of floats. + """ import pandas as pd - return pd.DataFrame(self._drain_dicts()) + columns = self._columns(25_000, None, True) + 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, 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), built in Python from the rows: a columnar JSON path, no - binary buffer. 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; 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: - 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, 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 58b261ba..778305b5 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], "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,42 @@ 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) -> 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). ``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) + 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: + 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..369c1dfa 100644 --- a/src/arcadedb/results.py +++ b/src/arcadedb/results.py @@ -12,6 +12,15 @@ import jpype as _jpype from jpype import JException +from ._columnar import ( + decode_arrow_frame, + decode_columns_frame, + empty_arrow, + empty_columns, + frame_rows, + merge_arrow, + merge_columns, +) from ._logging import get_logger from .callbacks import raise_pending_interrupt from .exceptions import ArcadeDBError @@ -92,6 +101,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 +170,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. @@ -691,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) and None for a null; + ``to_list()`` gives lists of floats. Args: convert_types: Convert Java types to Python (default: True) @@ -726,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)) @@ -744,7 +601,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: @@ -790,7 +649,13 @@ 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. 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 first appearance, because a document is schemaless: a property the @@ -816,142 +681,48 @@ 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 + for name in parts: + if name not in names: + names.append(name) + batches.append((count, parts)) + + return merge_columns(names, batches, columns) def to_arrow( self, batch_size: int = 25_000, columns: Optional[Sequence[str]] = None @@ -983,13 +754,11 @@ def to_arrow( falls back to the row-based paths. """ try: - import numpy as np - import pyarrow as pa + __import__("numpy") # availability check + __import__("pyarrow") # availability check except ImportError: return None - import json - column_batcher = _bridge_class("ColumnBatcher") if column_batcher is None: return None @@ -1000,141 +769,30 @@ 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=" 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/conftest.py b/tests/conftest.py index ec3147c7..1fdd6893 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -39,6 +39,8 @@ if sys.platform == "win32": collect_ignore.extend(WINDOWS_NOT_APPLICABLE) if BACKEND != "native": + # the old JSON-row route that the native column frames are compared with (the module lives in the native package) + collect_ignore.append("test_native_column_frames_reference.py") # Java -> 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..ff07ed59 --- /dev/null +++ b/tests/test_column_frames.py @@ -0,0 +1,607 @@ +"""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 pandas as pd +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] + # 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) + 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.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", "`` (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, + "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", + "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"), + ("bt", "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 _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): + 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() + 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 + else: + assert have.type == want.type, (name, have.type, want.type) + assert _eq(have.to_pylist(), 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 + np.testing.assert_array_equal(have, want, err_msg=name) + else: + assert isinstance(have, list), (name, type(have)) + 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): + 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: + 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: + # a null vector is a missing value, never a vector of NaN + assert cell is None, (name, cell) + 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 = [ + ("full", "SELECT FROM R ORDER BY big", None), + ( + "projection", + "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), + ("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_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") + 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()