vectrojs 0.0.0-stage → 0.1.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/CHANGELOG.md +71 -0
- package/LICENSE +21 -0
- package/README.md +822 -2
- package/bin/vectro.mjs +82 -0
- package/dist/vectro.wasm +0 -0
- package/mcp/server.mjs +288 -0
- package/package.json +91 -4
- package/src/arrow.mjs +905 -0
- package/src/bitpack.mjs +135 -0
- package/src/browser.mjs +129 -0
- package/src/bytes.mjs +60 -0
- package/src/csv-kernel.mjs +205 -0
- package/src/csv.mjs +270 -0
- package/src/dataset.mjs +1763 -0
- package/src/dates.mjs +139 -0
- package/src/emitter.mjs +109 -0
- package/src/engine.mjs +881 -0
- package/src/files-browser.mjs +104 -0
- package/src/files-node.mjs +59 -0
- package/src/frames.mjs +191 -0
- package/src/fuse.mjs +512 -0
- package/src/interp.mjs +305 -0
- package/src/join.mjs +534 -0
- package/src/kernels.mjs +312 -0
- package/src/ndjson.mjs +136 -0
- package/src/node.mjs +120 -0
- package/src/pack.mjs +114 -0
- package/src/parallel-sort.mjs +99 -0
- package/src/partition-kernels.mjs +388 -0
- package/src/partition.mjs +395 -0
- package/src/plan.mjs +551 -0
- package/src/portable.mjs +21 -0
- package/src/rank-kernels.mjs +81 -0
- package/src/rows.mjs +91 -0
- package/src/serve/catalog.mjs +148 -0
- package/src/serve/flight.mjs +270 -0
- package/src/serve/formats.mjs +87 -0
- package/src/serve/http.mjs +182 -0
- package/src/serve/index.mjs +9 -0
- package/src/serve/odata.mjs +291 -0
- package/src/serve/pg.mjs +588 -0
- package/src/serve/pgcatalog.mjs +220 -0
- package/src/serve/pgtypes.mjs +199 -0
- package/src/sketch.mjs +75 -0
- package/src/sort-kernels.mjs +117 -0
- package/src/sort.mjs +102 -0
- package/src/split-worker-core.mjs +252 -0
- package/src/split-worker-node.mjs +8 -0
- package/src/split-worker.mjs +7 -0
- package/src/split.mjs +741 -0
- package/src/sql.mjs +1558 -0
- package/src/sqltext.mjs +132 -0
- package/src/static-kernels.mjs +164 -0
- package/src/stream-worker.mjs +553 -0
- package/src/table-kernels.mjs +323 -0
- package/src/table.mjs +2206 -0
- package/src/textmap.mjs +159 -0
- package/src/vectors.mjs +138 -0
- package/src/wasm.mjs +254 -0
- package/src/window.mjs +160 -0
- package/src/worker-browser.mjs +21 -0
- package/src/worker-loop.mjs +106 -0
- package/src/worker.mjs +22 -0
- package/src/zones.mjs +75 -0
package/src/arrow.mjs
ADDED
|
@@ -0,0 +1,905 @@
|
|
|
1
|
+
// Apache Arrow IPC: the binary table format pandas, Polars, DuckDB, pyarrow
|
|
2
|
+
// and Arrow Flight use to move data. Reads the stream format and the file
|
|
3
|
+
// format (Feather v2), including LZ4-compressed bodies, and writes both.
|
|
4
|
+
//
|
|
5
|
+
// Arrow types map onto Vectro's two column kinds:
|
|
6
|
+
// number: ints, floats, bools (0/1), dates, times, timestamps and
|
|
7
|
+
// durations (as milliseconds); null -> NaN
|
|
8
|
+
// category: strings and dictionary-encoded strings; null -> ''
|
|
9
|
+
// Other types (lists, structs, decimals, binary...) are refused unless left
|
|
10
|
+
// out with `columns`. 64-bit integers are exact up to 2^53.
|
|
11
|
+
//
|
|
12
|
+
// Layout of the stream format: messages, each
|
|
13
|
+
// 0xFFFFFFFF, int32 metadata length, flatbuffer metadata (padded to 8), body
|
|
14
|
+
// then 0xFFFFFFFF 0x00000000 at the end. The file format wraps the stream in
|
|
15
|
+
// "ARROW1" magic and adds a footer.
|
|
16
|
+
|
|
17
|
+
import { ByteDict } from './csv-kernel.mjs';
|
|
18
|
+
|
|
19
|
+
const enc = new TextEncoder(), dec = new TextDecoder();
|
|
20
|
+
const pad8 = (n) => (n + 7) & ~7;
|
|
21
|
+
const MAGIC = enc.encode('ARROW1');
|
|
22
|
+
const isMagic = (b, at) => MAGIC.every((x, i) => b[at + i] === x);
|
|
23
|
+
const EOS = Uint8Array.of(255, 255, 255, 255, 0, 0, 0, 0);
|
|
24
|
+
const MESSAGE = { schema: 1, dictionary: 2, batch: 3 };
|
|
25
|
+
const V5 = 4;
|
|
26
|
+
|
|
27
|
+
// Arrow type ids (Schema.fbs, union Type)
|
|
28
|
+
const T = { Null: 1, Int: 2, FloatingPoint: 3, Binary: 4, Utf8: 5, Bool: 6, Decimal: 7, Date: 8, Time: 9, Timestamp: 10, Interval: 11, List: 12, Struct: 13, Union: 14, FixedSizeBinary: 15, FixedSizeList: 16, Map: 17, Duration: 18, LargeBinary: 19, LargeUtf8: 20, LargeList: 21, RunEndEncoded: 22, BinaryView: 23, Utf8View: 24, ListView: 25, LargeListView: 26 };
|
|
29
|
+
const TYPE_NAMES = Object.fromEntries(Object.entries(T).map(([k, v]) => [v, k]));
|
|
30
|
+
const NUMBER_TYPES = new Set([T.Null, T.Int, T.FloatingPoint, T.Bool, T.Date, T.Time, T.Timestamp, T.Duration]);
|
|
31
|
+
const TO_MS = [[1000, 1], [1, 1], [1, 1000], [1, 1e6]]; // TimeUnit s, ms, us, ns -> [multiply, divide]
|
|
32
|
+
|
|
33
|
+
// ---- flatbuffers, reading -----------------------------------------------------
|
|
34
|
+
|
|
35
|
+
class FlatReader {
|
|
36
|
+
constructor(bytes) {
|
|
37
|
+
this.b = bytes;
|
|
38
|
+
this.dv = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
|
|
39
|
+
}
|
|
40
|
+
root() { return this.dv.getUint32(0, true); }
|
|
41
|
+
#field(t, slot) {
|
|
42
|
+
const vt = t - this.dv.getInt32(t, true);
|
|
43
|
+
const at = 4 + 2 * slot;
|
|
44
|
+
if (at >= this.dv.getUint16(vt, true)) return 0;
|
|
45
|
+
const off = this.dv.getUint16(vt + at, true);
|
|
46
|
+
return off ? t + off : 0;
|
|
47
|
+
}
|
|
48
|
+
u8(t, slot, d = 0) { const p = this.#field(t, slot); return p ? this.b[p] : d; }
|
|
49
|
+
i16(t, slot, d = 0) { const p = this.#field(t, slot); return p ? this.dv.getInt16(p, true) : d; }
|
|
50
|
+
i32(t, slot, d = 0) { const p = this.#field(t, slot); return p ? this.dv.getInt32(p, true) : d; }
|
|
51
|
+
i64(t, slot, d = 0) { const p = this.#field(t, slot); return p ? this.long(p) : d; }
|
|
52
|
+
long(p) { return this.dv.getUint32(p, true) + this.dv.getInt32(p + 4, true) * 4294967296; }
|
|
53
|
+
table(t, slot) { const p = this.#field(t, slot); return p ? p + this.dv.getUint32(p, true) : 0; }
|
|
54
|
+
vector(t, slot) {
|
|
55
|
+
const p = this.#field(t, slot);
|
|
56
|
+
if (!p) return { at: 0, n: 0 };
|
|
57
|
+
const v = p + this.dv.getUint32(p, true);
|
|
58
|
+
return { at: v + 4, n: this.dv.getUint32(v, true) };
|
|
59
|
+
}
|
|
60
|
+
string(t, slot) { const v = this.vector(t, slot); return v.n ? dec.decode(this.b.subarray(v.at, v.at + v.n)) : ''; }
|
|
61
|
+
tables(t, slot) {
|
|
62
|
+
const v = this.vector(t, slot);
|
|
63
|
+
return Array.from({ length: v.n }, (_, i) => { const p = v.at + 4 * i; return p + this.dv.getUint32(p, true); });
|
|
64
|
+
}
|
|
65
|
+
}
|
|
66
|
+
|
|
67
|
+
function readField(fb, f) {
|
|
68
|
+
const typeId = fb.u8(f, 2), type = fb.table(f, 3), dict = fb.table(f, 4);
|
|
69
|
+
const field = { name: fb.string(f, 0), typeId, children: fb.tables(f, 5).map((c) => readField(fb, c)), dict: null };
|
|
70
|
+
if (typeId === T.Int) { field.bits = fb.i32(type, 0); field.signed = !!fb.u8(type, 1); }
|
|
71
|
+
else if (typeId === T.FloatingPoint) field.precision = fb.i16(type, 0);
|
|
72
|
+
else if (typeId === T.Date) field.unit = fb.i16(type, 0, 1);
|
|
73
|
+
else if (typeId === T.Time) { field.unit = fb.i16(type, 0, 1); field.bits = fb.i32(type, 1, 32); }
|
|
74
|
+
else if (typeId === T.Timestamp) field.unit = fb.i16(type, 0, 0);
|
|
75
|
+
else if (typeId === T.Duration) field.unit = fb.i16(type, 0, 1);
|
|
76
|
+
else if (typeId === T.Union) field.mode = fb.i16(type, 0, 0);
|
|
77
|
+
if (dict) {
|
|
78
|
+
const index = fb.table(dict, 1); // absent: int32
|
|
79
|
+
field.dict = { id: fb.i64(dict, 0), bits: index ? fb.i32(index, 0) : 32, signed: index ? !!fb.u8(index, 1) : true };
|
|
80
|
+
}
|
|
81
|
+
return field;
|
|
82
|
+
}
|
|
83
|
+
|
|
84
|
+
const valueKind = (t) => (t === T.Utf8 || t === T.LargeUtf8 ? 'category' : NUMBER_TYPES.has(t) ? 'number' : null);
|
|
85
|
+
const kindOf = (f) => (f.typeId === T.FloatingPoint && f.precision === 0 && !f.dict ? 'number' : valueKind(f.typeId));
|
|
86
|
+
const typeName = (f) => (f.dict ? `dictionary of ${TYPE_NAMES[f.typeId] ?? f.typeId}` : TYPE_NAMES[f.typeId] ?? `type ${f.typeId}`);
|
|
87
|
+
|
|
88
|
+
// A record batch's nodes (rows, nulls per column) and buffers (decompressed
|
|
89
|
+
// on demand).
|
|
90
|
+
function batchMeta(fb, rb, body) {
|
|
91
|
+
const nodes = fb.vector(rb, 1), buffers = fb.vector(rb, 2), comp = fb.table(rb, 3), variadic = fb.vector(rb, 4);
|
|
92
|
+
const codec = comp ? fb.u8(comp, 0, 0) : -1; // 0 = LZ4 frame, 1 = ZSTD
|
|
93
|
+
return {
|
|
94
|
+
length: fb.i64(rb, 0),
|
|
95
|
+
node: (i) => [fb.long(nodes.at + 16 * i), fb.long(nodes.at + 16 * i + 8)],
|
|
96
|
+
variadic: (i) => fb.long(variadic.at + 8 * i),
|
|
97
|
+
buffer(i) {
|
|
98
|
+
if (i >= buffers.n) throw new Error('vectro: an Arrow record batch has fewer buffers than its schema needs');
|
|
99
|
+
const off = fb.long(buffers.at + 16 * i), len = fb.long(buffers.at + 16 * i + 8);
|
|
100
|
+
if (off + len > body.length) throw new Error('vectro: an Arrow buffer runs past the end of its message (truncated data?)');
|
|
101
|
+
const raw = body.subarray(off, off + len);
|
|
102
|
+
if (codec < 0 || len === 0) return raw;
|
|
103
|
+
const dv = new DataView(raw.buffer, raw.byteOffset, 8);
|
|
104
|
+
const full = dv.getUint32(0, true) + dv.getInt32(4, true) * 4294967296;
|
|
105
|
+
if (full === -1) return raw.subarray(8); // stored uncompressed
|
|
106
|
+
if (codec !== 0) throw new Error('vectro: ZSTD-compressed Arrow data is not supported yet (write it uncompressed or with LZ4)');
|
|
107
|
+
return lz4Frame(raw.subarray(8), full);
|
|
108
|
+
},
|
|
109
|
+
};
|
|
110
|
+
}
|
|
111
|
+
|
|
112
|
+
function view(Type, bytes, n) {
|
|
113
|
+
const need = n * Type.BYTES_PER_ELEMENT;
|
|
114
|
+
if (bytes.length < need) throw new Error('vectro: an Arrow buffer is shorter than its column (truncated data?)');
|
|
115
|
+
if (bytes.byteOffset % Type.BYTES_PER_ELEMENT) bytes = bytes.slice(0, need); // misaligned: copy
|
|
116
|
+
return new Type(bytes.buffer, bytes.byteOffset, n);
|
|
117
|
+
}
|
|
118
|
+
|
|
119
|
+
const INT_TYPES = { 8: [Uint8Array, Int8Array], 16: [Uint16Array, Int16Array], 32: [Uint32Array, Int32Array] };
|
|
120
|
+
|
|
121
|
+
function ints(data, n, bits, signed, [mul, div]) {
|
|
122
|
+
let out;
|
|
123
|
+
if (bits === 64) {
|
|
124
|
+
out = new Float64Array(n);
|
|
125
|
+
const lo = view(Uint32Array, data, 2 * n), hi = signed ? view(Int32Array, data, 2 * n) : lo;
|
|
126
|
+
for (let i = 0; i < n; i++) out[i] = hi[2 * i + 1] * 4294967296 + lo[2 * i];
|
|
127
|
+
} else {
|
|
128
|
+
const Type = INT_TYPES[bits]?.[signed ? 1 : 0];
|
|
129
|
+
if (!Type) throw new Error(`vectro: unsupported Arrow integer width ${bits}`);
|
|
130
|
+
out = new Float64Array(n);
|
|
131
|
+
out.set(view(Type, data, n));
|
|
132
|
+
}
|
|
133
|
+
if (mul !== 1) for (let i = 0; i < n; i++) out[i] *= mul;
|
|
134
|
+
if (div !== 1) for (let i = 0; i < n; i++) out[i] /= div;
|
|
135
|
+
return out;
|
|
136
|
+
}
|
|
137
|
+
|
|
138
|
+
function half(h) {
|
|
139
|
+
const out = new Float64Array(h.length);
|
|
140
|
+
for (let i = 0; i < h.length; i++) {
|
|
141
|
+
const x = h[i], e = (x >> 10) & 31, m = x & 1023, s = x & 0x8000 ? -1 : 1;
|
|
142
|
+
out[i] = e === 0 ? s * m * 2 ** -24 : e === 31 ? (m ? NaN : s * Infinity) : s * (1 + m / 1024) * 2 ** (e - 15);
|
|
143
|
+
}
|
|
144
|
+
return out;
|
|
145
|
+
}
|
|
146
|
+
|
|
147
|
+
// Number column values (a view into `data` when it is float64 already).
|
|
148
|
+
function numbers(f, n, data) {
|
|
149
|
+
switch (f.typeId) {
|
|
150
|
+
case T.FloatingPoint:
|
|
151
|
+
if (f.precision === 2) return view(Float64Array, data, n);
|
|
152
|
+
if (f.precision === 1) { const out = new Float64Array(n); out.set(view(Float32Array, data, n)); return out; }
|
|
153
|
+
return half(view(Uint16Array, data, n));
|
|
154
|
+
case T.Int: return ints(data, n, f.bits, f.signed, [1, 1]);
|
|
155
|
+
case T.Bool: { const out = new Float64Array(n); for (let i = 0; i < n; i++) out[i] = (data[i >> 3] >> (i & 7)) & 1; return out; }
|
|
156
|
+
case T.Date: return f.unit === 0 ? ints(data, n, 32, true, [86_400_000, 1]) : ints(data, n, 64, true, [1, 1]);
|
|
157
|
+
case T.Time: return ints(data, n, f.bits, true, TO_MS[f.unit]);
|
|
158
|
+
default: return ints(data, n, 64, true, TO_MS[f.unit]); // Timestamp, Duration
|
|
159
|
+
}
|
|
160
|
+
}
|
|
161
|
+
|
|
162
|
+
const isNull = (valid, i) => !((valid[i >> 3] >> (i & 7)) & 1);
|
|
163
|
+
|
|
164
|
+
function indexArray(d, bytes, n) {
|
|
165
|
+
if (d.bits === 32) return view(Uint32Array, bytes, n); // non-negative int32: same bits
|
|
166
|
+
const out = new Uint32Array(n);
|
|
167
|
+
if (d.bits === 64) { const w = view(Uint32Array, bytes, 2 * n); for (let i = 0; i < n; i++) out[i] = w[2 * i]; return out; }
|
|
168
|
+
out.set(view(INT_TYPES[d.bits][d.signed ? 1 : 0], bytes, n));
|
|
169
|
+
return out;
|
|
170
|
+
}
|
|
171
|
+
|
|
172
|
+
// Strings as codes into the distinct values of this batch; a repeated value
|
|
173
|
+
// never becomes a JS string twice. Strings of up to 7 bytes (most labels) are
|
|
174
|
+
// packed into two ints and found in a small hash table; longer ones go
|
|
175
|
+
// through a byte-keyed dictionary.
|
|
176
|
+
const MASK_A = Int32Array.from({ length: 8 }, (_, n) => (n >= 4 ? -1 : (1 << (8 * n)) - 1));
|
|
177
|
+
const MASK_B = Int32Array.from({ length: 8 }, (_, n) => (n <= 4 ? 0 : (1 << (8 * (n - 4))) - 1));
|
|
178
|
+
|
|
179
|
+
function strings(offsets, data, n, large) {
|
|
180
|
+
const o = large ? view(Uint32Array, offsets, 2 * (n + 1)) : view(Int32Array, offsets, n + 1);
|
|
181
|
+
const step = large ? 2 : 1; // (large: the low half of each int64 offset)
|
|
182
|
+
const local = new Uint32Array(n), strs = [], longCodes = [];
|
|
183
|
+
const long = new ByteDict(data, (x) => x);
|
|
184
|
+
let cap = 256, lo = new Int32Array(cap), hi = new Int32Array(cap), code = new Int32Array(cap).fill(-1);
|
|
185
|
+
const slotOf = (a, b, mask) => {
|
|
186
|
+
const h = Math.imul(a, 0x9e3779b1) ^ Math.imul(b ^ (b >>> 16), 0x85ebca6b);
|
|
187
|
+
return (h ^ (h >>> 15)) & mask;
|
|
188
|
+
};
|
|
189
|
+
for (let i = 0; i < n; i++) {
|
|
190
|
+
const s = o[i * step], e = o[(i + 1) * step], len = e - s;
|
|
191
|
+
if (len > 7) {
|
|
192
|
+
const c = long.code(s, e, false);
|
|
193
|
+
let g = longCodes[c];
|
|
194
|
+
if (g === undefined) { g = longCodes[c] = strs.length; strs.push(long.strings[c]); }
|
|
195
|
+
local[i] = g;
|
|
196
|
+
continue;
|
|
197
|
+
}
|
|
198
|
+
// read 7 bytes regardless (no branch on the length), keep `len` of them
|
|
199
|
+
const a = (data[s] | (data[s + 1] << 8) | (data[s + 2] << 16) | (data[s + 3] << 24)) & MASK_A[len];
|
|
200
|
+
const b = (len << 24) | ((data[s + 4] | (data[s + 5] << 8) | (data[s + 6] << 16)) & MASK_B[len]);
|
|
201
|
+
const mask = cap - 1;
|
|
202
|
+
let slot = slotOf(a, b, mask);
|
|
203
|
+
while (code[slot] !== -1 && (lo[slot] !== a || hi[slot] !== b)) slot = (slot + 1) & mask;
|
|
204
|
+
if (code[slot] === -1) {
|
|
205
|
+
code[slot] = strs.length; lo[slot] = a; hi[slot] = b;
|
|
206
|
+
strs.push(dec.decode(data.subarray(s, e)));
|
|
207
|
+
if (strs.length * 2 > cap) { // grow and rehash
|
|
208
|
+
const old = [lo, hi, code];
|
|
209
|
+
cap *= 2; lo = new Int32Array(cap); hi = new Int32Array(cap); code = new Int32Array(cap).fill(-1);
|
|
210
|
+
for (let k = 0; k < old[2].length; k++) {
|
|
211
|
+
if (old[2][k] === -1) continue;
|
|
212
|
+
let t = slotOf(old[0][k], old[1][k], cap - 1);
|
|
213
|
+
while (code[t] !== -1) t = (t + 1) & (cap - 1);
|
|
214
|
+
lo[t] = old[0][k]; hi[t] = old[1][k]; code[t] = old[2][k];
|
|
215
|
+
}
|
|
216
|
+
local[i] = strs.length - 1;
|
|
217
|
+
continue;
|
|
218
|
+
}
|
|
219
|
+
}
|
|
220
|
+
local[i] = code[slot];
|
|
221
|
+
}
|
|
222
|
+
return { local, strings: strs };
|
|
223
|
+
}
|
|
224
|
+
|
|
225
|
+
// One top-level column: { values } | { local, strings, valid } | { indices, valid }.
|
|
226
|
+
function decodeColumn(f, meta, st) {
|
|
227
|
+
const [n, nulls] = meta.node(st.node++);
|
|
228
|
+
const take = () => meta.buffer(st.buf++);
|
|
229
|
+
if (f.dict) {
|
|
230
|
+
const valid = take(), idx = take();
|
|
231
|
+
return { n, valid: nulls && valid.length ? valid : null, indices: indexArray(f.dict, idx, n) };
|
|
232
|
+
}
|
|
233
|
+
if (f.typeId === T.Null) return { n, values: new Float64Array(n).fill(NaN) };
|
|
234
|
+
const valid = take();
|
|
235
|
+
if (f.typeId === T.Utf8 || f.typeId === T.LargeUtf8) {
|
|
236
|
+
const offsets = take(), data = take();
|
|
237
|
+
return { n, valid: nulls && valid.length ? valid : null, ...strings(offsets, data, n, f.typeId === T.LargeUtf8) };
|
|
238
|
+
}
|
|
239
|
+
const data = take();
|
|
240
|
+
let values = numbers(f, n, data);
|
|
241
|
+
if (nulls && valid.length) {
|
|
242
|
+
if (values.buffer === data.buffer) values = values.slice(); // don't write into the caller's bytes
|
|
243
|
+
for (let i = 0; i < n; i++) if (isNull(valid, i)) values[i] = NaN;
|
|
244
|
+
}
|
|
245
|
+
return { n, values };
|
|
246
|
+
}
|
|
247
|
+
|
|
248
|
+
// Step over a column that isn't read (any type, nested or not).
|
|
249
|
+
function skipColumn(f, meta, st) {
|
|
250
|
+
st.node++;
|
|
251
|
+
if (f.dict) { st.buf += 2; return; }
|
|
252
|
+
const t = f.typeId;
|
|
253
|
+
if (t === T.Null || t === T.RunEndEncoded) st.buf += 0;
|
|
254
|
+
else if (t === T.Struct || t === T.FixedSizeList) st.buf += 1;
|
|
255
|
+
else if (t === T.Union) st.buf += f.mode === 1 ? 2 : 1;
|
|
256
|
+
else if (t === T.Binary || t === T.Utf8 || t === T.LargeBinary || t === T.LargeUtf8 || t === T.ListView || t === T.LargeListView) st.buf += 3;
|
|
257
|
+
else if (t === T.BinaryView || t === T.Utf8View) st.buf += 2 + meta.variadic(st.view++);
|
|
258
|
+
else st.buf += 2;
|
|
259
|
+
for (const c of f.children) skipColumn(c, meta, st);
|
|
260
|
+
}
|
|
261
|
+
|
|
262
|
+
// ---- reading, incrementally -----------------------------------------------------
|
|
263
|
+
|
|
264
|
+
// Reads Arrow IPC bytes that arrive in pieces cut anywhere (the file format
|
|
265
|
+
// works too: its footer is ignored). Calls
|
|
266
|
+
// onSchema(kinds, fields): kinds = { name: 'number' | 'category' | null }
|
|
267
|
+
// (null: a type Vectro can't read); may return the names to read
|
|
268
|
+
// (default: every column it can read);
|
|
269
|
+
// onBatch(rows, columns): columns[k] (schema order) = { values: Float64Array }
|
|
270
|
+
// or { codes: Uint32Array, labels: string[] } (all labels of column k so
|
|
271
|
+
// far; codes stay valid as labels are added), or null if not read;
|
|
272
|
+
// onEnd() at the end-of-stream marker.
|
|
273
|
+
// Views may point into the bytes pushed: copy what you keep.
|
|
274
|
+
export class ArrowReader {
|
|
275
|
+
constructor({ onSchema, onBatch, onEnd = () => {} }) {
|
|
276
|
+
this.on = { onSchema, onBatch, onEnd };
|
|
277
|
+
this.buf = new Uint8Array(1 << 16);
|
|
278
|
+
this.len = 0;
|
|
279
|
+
this.started = false;
|
|
280
|
+
this.ended = false;
|
|
281
|
+
this.fields = null;
|
|
282
|
+
this.cols = null; // per field: { labels, index, lut, version } or null
|
|
283
|
+
this.dicts = new Map(); // dictionary id -> { field, values, numeric, version }
|
|
284
|
+
}
|
|
285
|
+
|
|
286
|
+
push(chunk) {
|
|
287
|
+
if (this.ended) return;
|
|
288
|
+
if (!this.len) { // nothing buffered: read straight from the chunk
|
|
289
|
+
const used = this.#scan(chunk, chunk.length);
|
|
290
|
+
this.#keep(chunk.subarray(used));
|
|
291
|
+
return;
|
|
292
|
+
}
|
|
293
|
+
this.#keep(chunk);
|
|
294
|
+
const used = this.#scan(this.buf, this.len);
|
|
295
|
+
this.buf.copyWithin(0, used, this.len);
|
|
296
|
+
this.len -= used;
|
|
297
|
+
}
|
|
298
|
+
|
|
299
|
+
#keep(bytes) {
|
|
300
|
+
if (!bytes.length || this.ended) return;
|
|
301
|
+
if (this.len + bytes.length > this.buf.length) {
|
|
302
|
+
const bigger = new Uint8Array(Math.max(this.buf.length * 2, this.len + bytes.length));
|
|
303
|
+
bigger.set(this.buf.subarray(0, this.len));
|
|
304
|
+
this.buf = bigger;
|
|
305
|
+
}
|
|
306
|
+
this.buf.set(bytes, this.len);
|
|
307
|
+
this.len += bytes.length;
|
|
308
|
+
}
|
|
309
|
+
|
|
310
|
+
#scan(b, len) {
|
|
311
|
+
const dv = new DataView(b.buffer, b.byteOffset, b.byteLength);
|
|
312
|
+
let at = 0;
|
|
313
|
+
if (!this.started) {
|
|
314
|
+
if (len < 8) return 0;
|
|
315
|
+
this.started = true;
|
|
316
|
+
if (isMagic(b, 0)) at = 8; // file format
|
|
317
|
+
}
|
|
318
|
+
while (len - at >= 4) {
|
|
319
|
+
let p = at, size = dv.getInt32(p, true);
|
|
320
|
+
if (size === -1) { if (len - at < 8) break; size = dv.getInt32(p + 4, true); p += 8; } else p += 4; // (old streams have no 0xFFFFFFFF)
|
|
321
|
+
if (size === 0) { at = p; this.ended = true; this.on.onEnd(); break; }
|
|
322
|
+
if (size < 0) throw new Error('vectro: this is not Arrow IPC data (bad message length)');
|
|
323
|
+
if (len - p < size) break;
|
|
324
|
+
const fb = new FlatReader(b.subarray(p, p + size));
|
|
325
|
+
const msg = fb.root();
|
|
326
|
+
const bodyLength = fb.i64(msg, 3);
|
|
327
|
+
if (len - p - size < bodyLength) break;
|
|
328
|
+
this.#message(fb, msg, b.subarray(p + size, p + size + bodyLength));
|
|
329
|
+
at = p + size + bodyLength;
|
|
330
|
+
}
|
|
331
|
+
return at;
|
|
332
|
+
}
|
|
333
|
+
|
|
334
|
+
#message(fb, msg, body) {
|
|
335
|
+
const version = fb.i16(msg, 0);
|
|
336
|
+
if (version < 3) throw new Error(`vectro: Arrow IPC metadata V${version + 1} is too old (V4 or later)`);
|
|
337
|
+
const kind = fb.u8(msg, 1), h = fb.table(msg, 2);
|
|
338
|
+
if (kind === MESSAGE.schema) return this.#schema(fb, h);
|
|
339
|
+
if (!this.fields) throw new Error('vectro: Arrow data arrived before the schema');
|
|
340
|
+
if (kind === MESSAGE.dictionary) return this.#dictionary(fb, h, body);
|
|
341
|
+
if (kind === MESSAGE.batch) return this.#batch(fb, h, body);
|
|
342
|
+
throw new Error(`vectro: unsupported Arrow message type ${kind}`);
|
|
343
|
+
}
|
|
344
|
+
|
|
345
|
+
#schema(fb, h) {
|
|
346
|
+
if (this.fields) throw new Error('vectro: a second Arrow schema arrived on one stream');
|
|
347
|
+
if (fb.i16(h, 0) === 1) throw new Error('vectro: big-endian Arrow data is not supported');
|
|
348
|
+
this.fields = fb.tables(h, 1).map((f) => readField(fb, f));
|
|
349
|
+
const kinds = Object.fromEntries(this.fields.map((f) => [f.name, kindOf(f)]));
|
|
350
|
+
const info = this.fields.map((f) => ({ name: f.name, kind: kindOf(f), type: typeName(f), date: !f.dict && (f.typeId === T.Date || f.typeId === T.Timestamp) }));
|
|
351
|
+
const want = new Set(this.on.onSchema(kinds, info) ?? info.filter((f) => f.kind).map((f) => f.name));
|
|
352
|
+
this.cols = this.fields.map((f) => (want.has(f.name) && kindOf(f) ? { labels: [], index: new Map(), lut: [], identity: true, version: -1 } : null));
|
|
353
|
+
this.fields.forEach((f, k) => {
|
|
354
|
+
if (!this.cols[k] || !f.dict) return;
|
|
355
|
+
if (!this.dicts.has(f.dict.id)) this.dicts.set(f.dict.id, { field: { ...f, dict: null }, values: null, numeric: kindOf(f) === 'number', version: 0 });
|
|
356
|
+
});
|
|
357
|
+
}
|
|
358
|
+
|
|
359
|
+
#dictionary(fb, h, body) {
|
|
360
|
+
const d = this.dicts.get(fb.i64(h, 0));
|
|
361
|
+
if (!d) return; // a dictionary for a column not read
|
|
362
|
+
const c = decodeColumn(d.field, batchMeta(fb, fb.table(h, 1), body), { node: 0, buf: 0, view: 0 });
|
|
363
|
+
const delta = !!fb.u8(h, 2);
|
|
364
|
+
if (!delta || !d.values) { d.values = []; d.version++; }
|
|
365
|
+
if (c.values) { for (const v of c.values) d.values.push(v); return; }
|
|
366
|
+
for (let i = 0; i < c.n; i++) d.values.push(c.valid && isNull(c.valid, i) ? '' : c.strings[c.local[i]]);
|
|
367
|
+
}
|
|
368
|
+
|
|
369
|
+
#label(col, s) {
|
|
370
|
+
let g = col.index.get(s);
|
|
371
|
+
if (g === undefined) { g = col.labels.length; col.labels.push(s); col.index.set(s, g); }
|
|
372
|
+
return g;
|
|
373
|
+
}
|
|
374
|
+
|
|
375
|
+
#batch(fb, h, body) {
|
|
376
|
+
const meta = batchMeta(fb, h, body);
|
|
377
|
+
const st = { node: 0, buf: 0, view: 0 };
|
|
378
|
+
const out = this.fields.map((f, k) => {
|
|
379
|
+
const col = this.cols[k];
|
|
380
|
+
if (!col) { skipColumn(f, meta, st); return null; }
|
|
381
|
+
const c = decodeColumn(f, meta, st);
|
|
382
|
+
if (c.values) return { values: c.values };
|
|
383
|
+
if (c.indices) return this.#fromDictionary(f, col, c);
|
|
384
|
+
const lut = c.strings.map((s) => this.#label(col, s));
|
|
385
|
+
const codes = new Uint32Array(c.n);
|
|
386
|
+
for (let i = 0; i < c.n; i++) codes[i] = lut[c.local[i]];
|
|
387
|
+
if (c.valid) { const e = this.#label(col, ''); for (let i = 0; i < c.n; i++) if (isNull(c.valid, i)) codes[i] = e; }
|
|
388
|
+
return { codes, labels: col.labels };
|
|
389
|
+
});
|
|
390
|
+
this.on.onBatch(meta.length, out);
|
|
391
|
+
}
|
|
392
|
+
|
|
393
|
+
#fromDictionary(f, col, c) {
|
|
394
|
+
const d = this.dicts.get(f.dict.id), idx = c.indices, n = c.n;
|
|
395
|
+
if (!d?.values) throw new Error(`vectro: Arrow column "${f.name}" arrived before its dictionary`);
|
|
396
|
+
if (d.numeric) {
|
|
397
|
+
const values = new Float64Array(n);
|
|
398
|
+
for (let i = 0; i < n; i++) values[i] = c.valid && isNull(c.valid, i) ? NaN : d.values[idx[i]];
|
|
399
|
+
return { values };
|
|
400
|
+
}
|
|
401
|
+
if (col.version !== d.version) { col.version = d.version; col.lut = []; col.identity = true; } // dictionary replaced
|
|
402
|
+
for (let j = col.lut.length; j < d.values.length; j++) {
|
|
403
|
+
const g = this.#label(col, d.values[j]);
|
|
404
|
+
col.lut.push(g);
|
|
405
|
+
if (g !== j) col.identity = false;
|
|
406
|
+
}
|
|
407
|
+
let codes = idx;
|
|
408
|
+
if (!col.identity || !(idx instanceof Uint32Array)) {
|
|
409
|
+
const lut = col.lut;
|
|
410
|
+
codes = new Uint32Array(n);
|
|
411
|
+
for (let i = 0; i < n; i++) codes[i] = lut[idx[i]];
|
|
412
|
+
}
|
|
413
|
+
if (c.valid) {
|
|
414
|
+
const e = this.#label(col, '');
|
|
415
|
+
if (codes === idx) codes = codes.slice();
|
|
416
|
+
for (let i = 0; i < n; i++) if (isNull(c.valid, i)) codes[i] = e;
|
|
417
|
+
}
|
|
418
|
+
return { codes, labels: col.labels };
|
|
419
|
+
}
|
|
420
|
+
}
|
|
421
|
+
|
|
422
|
+
// Read a whole Arrow stream or file. Returns { rows, columns, dates } with
|
|
423
|
+
// columns ready for engine.table(): { name: Float64Array | { codes, labels } },
|
|
424
|
+
// and the names of the date and timestamp columns (values in ms).
|
|
425
|
+
// columns: the names to read (default: all; a type Vectro can't read is an
|
|
426
|
+
// error unless left out).
|
|
427
|
+
export function readArrow(bytes, { columns: pick } = {}) {
|
|
428
|
+
if (!(bytes instanceof Uint8Array)) bytes = ArrayBuffer.isView(bytes) ? new Uint8Array(bytes.buffer, bytes.byteOffset, bytes.byteLength) : new Uint8Array(bytes);
|
|
429
|
+
let names = null, dates = [];
|
|
430
|
+
const parts = [];
|
|
431
|
+
const reader = new ArrowReader({
|
|
432
|
+
onSchema(kinds, info) {
|
|
433
|
+
const want = pick ?? Object.keys(kinds);
|
|
434
|
+
for (const name of want) {
|
|
435
|
+
const f = info.find((x) => x.name === name);
|
|
436
|
+
if (!f) throw new Error(`vectro: the Arrow data has no column "${name}" (columns: ${info.map((x) => x.name).join(', ')})`);
|
|
437
|
+
if (!f.kind) throw new Error(`vectro: column "${name}" has Arrow type ${f.type}, which Vectro can't read yet; leave it out with columns: [...]`);
|
|
438
|
+
}
|
|
439
|
+
names = info.map((x) => x.name);
|
|
440
|
+
dates = info.filter((x) => x.date && want.includes(x.name)).map((x) => x.name);
|
|
441
|
+
return want;
|
|
442
|
+
},
|
|
443
|
+
onBatch(rows, cols) { parts.push([rows, cols]); }, // views stay valid: all bytes are pushed at once
|
|
444
|
+
});
|
|
445
|
+
reader.push(bytes);
|
|
446
|
+
if (!names) throw new Error('vectro: no Arrow schema found (is this Arrow IPC data?)');
|
|
447
|
+
if (reader.len && !reader.ended) throw new Error('vectro: the Arrow data ends in the middle of a message (truncated?)');
|
|
448
|
+
const rows = parts.reduce((a, [n]) => a + n, 0);
|
|
449
|
+
const columns = {};
|
|
450
|
+
const want = pick ?? names;
|
|
451
|
+
for (const name of want) {
|
|
452
|
+
const k = names.indexOf(name);
|
|
453
|
+
const pieces = parts.map(([, cols]) => cols[k]);
|
|
454
|
+
const first = pieces[0];
|
|
455
|
+
if (!first) { columns[name] = new Float64Array(0); continue; }
|
|
456
|
+
if (first.values) {
|
|
457
|
+
columns[name] = pieces.length === 1 ? first.values : concatTyped(Float64Array, pieces.map((p) => p.values), rows);
|
|
458
|
+
} else {
|
|
459
|
+
const codes = pieces.length === 1 ? first.codes : concatTyped(Uint32Array, pieces.map((p) => p.codes), rows);
|
|
460
|
+
columns[name] = { codes, labels: pieces[pieces.length - 1].labels };
|
|
461
|
+
}
|
|
462
|
+
}
|
|
463
|
+
return { rows, columns, dates };
|
|
464
|
+
}
|
|
465
|
+
|
|
466
|
+
function concatTyped(Type, list, n) {
|
|
467
|
+
const out = new Type(n);
|
|
468
|
+
let at = 0;
|
|
469
|
+
for (const a of list) { out.set(a, at); at += a.length; }
|
|
470
|
+
return out;
|
|
471
|
+
}
|
|
472
|
+
|
|
473
|
+
// ---- LZ4 (frame format), for compressed Arrow bodies -----------------------------
|
|
474
|
+
|
|
475
|
+
function lz4Frame(src, size) {
|
|
476
|
+
const dv = new DataView(src.buffer, src.byteOffset, src.byteLength);
|
|
477
|
+
if (dv.getUint32(0, true) !== 0x184d2204) throw new Error('vectro: bad LZ4 frame in Arrow data');
|
|
478
|
+
const out = new Uint8Array(size);
|
|
479
|
+
const flg = src[4];
|
|
480
|
+
let p = 6 + (flg & 8 ? 8 : 0) + (flg & 1 ? 4 : 0) + 1; // header: magic, FLG, BD, [size], [dict id], checksum
|
|
481
|
+
let o = 0;
|
|
482
|
+
for (;;) {
|
|
483
|
+
const word = dv.getUint32(p, true);
|
|
484
|
+
p += 4;
|
|
485
|
+
if (word === 0) break;
|
|
486
|
+
const len = word & 0x7fffffff;
|
|
487
|
+
if (word & 0x80000000) { out.set(src.subarray(p, p + len), o); o += len; }
|
|
488
|
+
else o = lz4Block(src, p, p + len, out, o);
|
|
489
|
+
p += len + (flg & 0x10 ? 4 : 0);
|
|
490
|
+
}
|
|
491
|
+
if (o !== size) throw new Error('vectro: an LZ4 buffer in the Arrow data has the wrong size');
|
|
492
|
+
return out;
|
|
493
|
+
}
|
|
494
|
+
|
|
495
|
+
function lz4Block(src, p, end, out, o) {
|
|
496
|
+
while (p < end) {
|
|
497
|
+
const token = src[p++];
|
|
498
|
+
let lit = token >> 4;
|
|
499
|
+
if (lit === 15) { let b; do { b = src[p++]; lit += b; } while (b === 255); }
|
|
500
|
+
if (lit < 32) for (let i = 0; i < lit; i++) out[o++] = src[p++];
|
|
501
|
+
else { out.set(src.subarray(p, p + lit), o); o += lit; p += lit; }
|
|
502
|
+
if (p >= end) break;
|
|
503
|
+
const off = src[p] | (src[p + 1] << 8);
|
|
504
|
+
p += 2;
|
|
505
|
+
let m = token & 15;
|
|
506
|
+
if (m === 15) { let b; do { b = src[p++]; m += b; } while (b === 255); }
|
|
507
|
+
m += 4;
|
|
508
|
+
let from = o - off;
|
|
509
|
+
if (m < 32) for (let i = 0; i < m; i++) out[o++] = out[from++]; // short (or overlapping) copy
|
|
510
|
+
else { // long: the repeating part doubles with each copy
|
|
511
|
+
const stop = o + m;
|
|
512
|
+
while (o < stop) { const k = Math.min(stop - o, o - from); out.copyWithin(o, from, from + k); o += k; }
|
|
513
|
+
}
|
|
514
|
+
}
|
|
515
|
+
return o;
|
|
516
|
+
}
|
|
517
|
+
|
|
518
|
+
// ---- flatbuffers, writing -----------------------------------------------------------
|
|
519
|
+
// Objects are written front to back (a table, then what it points to), so
|
|
520
|
+
// every offset points forward as flatbuffers require.
|
|
521
|
+
// A table is { table: [slot, ...] }, each slot null or [kind, value] with kind
|
|
522
|
+
// u8 | bool | i16 | i32 | i64 | off (value: a table, string, { vec: [...] } of
|
|
523
|
+
// tables, or { structs: bytes, count } with 8-byte-aligned elements).
|
|
524
|
+
|
|
525
|
+
const SIZE = { u8: 1, bool: 1, i16: 2, i32: 4, off: 4, i64: 8 };
|
|
526
|
+
const tbl = (...slots) => ({ table: slots });
|
|
527
|
+
|
|
528
|
+
class FlatWriter {
|
|
529
|
+
constructor() { this.b = new Uint8Array(512); this.dv = new DataView(this.b.buffer); this.n = 0; }
|
|
530
|
+
#room(k) {
|
|
531
|
+
if (this.n + k <= this.b.length) return;
|
|
532
|
+
const bigger = new Uint8Array(Math.max(this.b.length * 2, this.n + k));
|
|
533
|
+
bigger.set(this.b);
|
|
534
|
+
this.b = bigger;
|
|
535
|
+
this.dv = new DataView(bigger.buffer);
|
|
536
|
+
}
|
|
537
|
+
#align(a) { const k = (a - (this.n % a)) % a; this.#room(k); this.n += k; }
|
|
538
|
+
#u16(v) { this.#room(2); this.dv.setUint16(this.n, v, true); this.n += 2; }
|
|
539
|
+
#u32(v) { this.#room(4); this.dv.setUint32(this.n, v, true); this.n += 4; }
|
|
540
|
+
#bytes(x) { this.#room(x.length); this.b.set(x, this.n); this.n += x.length; }
|
|
541
|
+
finish(root) {
|
|
542
|
+
this.#u32(0);
|
|
543
|
+
const at = this.#write(root); // (may grow the buffer: write the offset after)
|
|
544
|
+
this.dv.setUint32(0, at, true);
|
|
545
|
+
this.#align(8);
|
|
546
|
+
return this.b.slice(0, this.n);
|
|
547
|
+
}
|
|
548
|
+
#write(x) {
|
|
549
|
+
if (typeof x === 'string') {
|
|
550
|
+
this.#align(4);
|
|
551
|
+
const at = this.n, s = enc.encode(x);
|
|
552
|
+
this.#u32(s.length); this.#bytes(s); this.#bytes([0]);
|
|
553
|
+
return at;
|
|
554
|
+
}
|
|
555
|
+
if (x.vec) {
|
|
556
|
+
this.#align(4);
|
|
557
|
+
const at = this.n;
|
|
558
|
+
this.#u32(x.vec.length);
|
|
559
|
+
const slots = x.vec.map(() => { const p = this.n; this.#u32(0); return p; });
|
|
560
|
+
x.vec.forEach((child, i) => { const c = this.#write(child); this.dv.setUint32(slots[i], c - slots[i], true); });
|
|
561
|
+
return at;
|
|
562
|
+
}
|
|
563
|
+
if (x.structs) {
|
|
564
|
+
this.#align(4);
|
|
565
|
+
if ((this.n + 4) % 8) this.#u32(0);
|
|
566
|
+
const at = this.n;
|
|
567
|
+
this.#u32(x.count); this.#bytes(x.structs);
|
|
568
|
+
return at;
|
|
569
|
+
}
|
|
570
|
+
const slots = x.table;
|
|
571
|
+
const order = slots.map((s, i) => [i, s]).filter(([, s]) => s).sort((a, b) => SIZE[b[1][0]] - SIZE[a[1][0]]);
|
|
572
|
+
const pos = [];
|
|
573
|
+
let cur = 4; // after the vtable offset
|
|
574
|
+
for (const [i, [kind]] of order) { cur = (cur + SIZE[kind] - 1) & -SIZE[kind]; pos[i] = cur; cur += SIZE[kind]; }
|
|
575
|
+
this.#align(2);
|
|
576
|
+
const vt = this.n;
|
|
577
|
+
this.#u16(4 + 2 * slots.length); this.#u16(cur);
|
|
578
|
+
for (let i = 0; i < slots.length; i++) this.#u16(slots[i] ? pos[i] : 0);
|
|
579
|
+
this.#align(order.some(([, [k]]) => k === 'i64') ? 8 : 4);
|
|
580
|
+
const t = this.n;
|
|
581
|
+
this.#room(cur);
|
|
582
|
+
this.n += cur;
|
|
583
|
+
this.dv.setInt32(t, t - vt, true);
|
|
584
|
+
const later = [];
|
|
585
|
+
for (const [i, [kind, v]] of order) {
|
|
586
|
+
const p = t + pos[i];
|
|
587
|
+
if (kind === 'u8' || kind === 'bool') this.b[p] = +v;
|
|
588
|
+
else if (kind === 'i16') this.dv.setInt16(p, v, true);
|
|
589
|
+
else if (kind === 'i32') this.dv.setInt32(p, v, true);
|
|
590
|
+
else if (kind === 'i64') this.dv.setBigInt64(p, BigInt(v), true);
|
|
591
|
+
else later.push([p, v]);
|
|
592
|
+
}
|
|
593
|
+
for (const [p, child] of later) { const c = this.#write(child); this.dv.setUint32(p, c - p, true); }
|
|
594
|
+
return t;
|
|
595
|
+
}
|
|
596
|
+
}
|
|
597
|
+
|
|
598
|
+
// ---- writing --------------------------------------------------------------------------
|
|
599
|
+
|
|
600
|
+
const KIND_TYPES = { number: [T.FloatingPoint, 2], float32: [T.FloatingPoint, 1], category: [T.Utf8], date: [T.Timestamp] };
|
|
601
|
+
|
|
602
|
+
function schemaTable(columns) {
|
|
603
|
+
const fields = columns.map(([name, kind], k) => {
|
|
604
|
+
const [typeId, precision] = KIND_TYPES[kind];
|
|
605
|
+
const dict = kind === 'category' ? ['off', tbl(['i64', k], ['off', tbl(['i32', 32], ['bool', 1])], ['bool', 0])] : null;
|
|
606
|
+
const type = kind === 'date' ? tbl(['i16', 1], ['off', 'UTC']) : precision ? tbl(['i16', precision]) : tbl(); // date: Timestamp(ms, UTC)
|
|
607
|
+
return tbl(['off', name], ['bool', 1], ['u8', typeId], ['off', type], dict, ['off', { vec: [] }]);
|
|
608
|
+
});
|
|
609
|
+
return tbl(['i16', 0], ['off', { vec: fields }]);
|
|
610
|
+
}
|
|
611
|
+
|
|
612
|
+
function longs(pairs) { // [[a, b], ...] as int64 structs
|
|
613
|
+
const dv = new DataView(new ArrayBuffer(16 * pairs.length));
|
|
614
|
+
pairs.forEach(([a, b], i) => { dv.setBigInt64(16 * i, BigInt(a), true); dv.setBigInt64(16 * i + 8, BigInt(b), true); });
|
|
615
|
+
return { structs: new Uint8Array(dv.buffer), count: pairs.length };
|
|
616
|
+
}
|
|
617
|
+
|
|
618
|
+
function recordBatch(rows, nodes, buffers) {
|
|
619
|
+
let at = 0;
|
|
620
|
+
const specs = buffers.map((b) => { const s = [at, b.length]; at += pad8(b.length); return s; });
|
|
621
|
+
return { header: tbl(['i64', rows], ['off', longs(nodes)], ['off', longs(specs)]), buffers, bodyLength: at };
|
|
622
|
+
}
|
|
623
|
+
|
|
624
|
+
// One message: { bytes, metaLength (with its 8-byte prefix), bodyLength }.
|
|
625
|
+
function message(type, header, buffers = [], bodyLength = 0) {
|
|
626
|
+
const meta = new FlatWriter().finish(tbl(['i16', V5], ['u8', type], ['off', header], ['i64', bodyLength]));
|
|
627
|
+
const bytes = alloc(8 + meta.length + bodyLength); // not zero-filled: every byte is written below
|
|
628
|
+
const dv = new DataView(bytes.buffer, bytes.byteOffset, bytes.length);
|
|
629
|
+
dv.setUint32(0, 0xffffffff, true);
|
|
630
|
+
dv.setInt32(4, meta.length, true);
|
|
631
|
+
bytes.set(meta, 8);
|
|
632
|
+
let at = 8 + meta.length;
|
|
633
|
+
for (const b of buffers) {
|
|
634
|
+
bytes.set(b, at);
|
|
635
|
+
const end = at + pad8(b.length);
|
|
636
|
+
bytes.fill(0, at + b.length, end);
|
|
637
|
+
at = end;
|
|
638
|
+
}
|
|
639
|
+
return { bytes, metaLength: 8 + meta.length, bodyLength };
|
|
640
|
+
}
|
|
641
|
+
|
|
642
|
+
// Node's Buffer.allocUnsafeSlow skips zero-filling; elsewhere a Uint8Array.
|
|
643
|
+
const alloc = globalThis.Buffer?.allocUnsafeSlow ? (n) => globalThis.Buffer.allocUnsafeSlow(n) : (n) => new Uint8Array(n);
|
|
644
|
+
const asBytes = (a) => new Uint8Array(a.buffer, a.byteOffset, a.byteLength);
|
|
645
|
+
|
|
646
|
+
// Builds Arrow IPC stream messages for one connection or file: number columns
|
|
647
|
+
// as float64 (or float32), category columns dictionary-encoded with int32
|
|
648
|
+
// indices; labels new since the last batch go out as a delta dictionary.
|
|
649
|
+
// columns: { name: 'number' | 'float32' | 'category' }.
|
|
650
|
+
export class ArrowEncoder {
|
|
651
|
+
constructor(columns) {
|
|
652
|
+
this.columns = Object.entries(columns);
|
|
653
|
+
for (const [name, kind] of this.columns) if (!KIND_TYPES[kind]) throw new Error(`vectro: Arrow column "${name}" must be 'number' or 'category', got "${kind}"`);
|
|
654
|
+
this.labels = this.columns.map(([, kind]) => (kind === 'category' ? new Map() : null));
|
|
655
|
+
this.sent = this.columns.map(() => false);
|
|
656
|
+
}
|
|
657
|
+
|
|
658
|
+
schema() { return message(MESSAGE.schema, schemaTable(this.columns)).bytes; }
|
|
659
|
+
end() { return EOS.slice(); }
|
|
660
|
+
encode(batch) {
|
|
661
|
+
const parts = this.messages(batch);
|
|
662
|
+
return parts.length === 1 ? parts[0].bytes : concatBytes(parts.map((p) => p.bytes));
|
|
663
|
+
}
|
|
664
|
+
|
|
665
|
+
// batch: { name: Float64Array | number[] | string[] | { codes, labels } }.
|
|
666
|
+
// Returns the messages: dictionaries first, then the record batch.
|
|
667
|
+
messages(batch) {
|
|
668
|
+
let rows = -1;
|
|
669
|
+
const out = [], nodes = [], buffers = [];
|
|
670
|
+
this.columns.forEach(([name, kind], k) => {
|
|
671
|
+
const values = batch[name];
|
|
672
|
+
if (!values) throw new Error(`vectro: batch is missing column "${name}"`);
|
|
673
|
+
const len = values.codes ? values.codes.length : values.length;
|
|
674
|
+
if (rows === -1) rows = len;
|
|
675
|
+
else if (len !== rows) throw new RangeError(`vectro: column "${name}" has ${len} values, expected ${rows}`);
|
|
676
|
+
if (kind === 'date') { // int64 ms, missing (NaN) as null
|
|
677
|
+
const words = new Int32Array(2 * len), valid = new Uint8Array((len + 7) >> 3);
|
|
678
|
+
let nulls = 0;
|
|
679
|
+
for (let i = 0; i < len; i++) {
|
|
680
|
+
const v = values[i];
|
|
681
|
+
if (v !== v || v == null) { nulls++; continue; }
|
|
682
|
+
const ms = Math.round(v), hi = Math.floor(ms / 4294967296);
|
|
683
|
+
words[2 * i] = ms - hi * 4294967296; words[2 * i + 1] = hi;
|
|
684
|
+
valid[i >> 3] |= 1 << (i & 7);
|
|
685
|
+
}
|
|
686
|
+
nodes.push([len, nulls]);
|
|
687
|
+
buffers.push(nulls ? valid : new Uint8Array(0), asBytes(words));
|
|
688
|
+
return;
|
|
689
|
+
}
|
|
690
|
+
nodes.push([len, 0]);
|
|
691
|
+
if (kind !== 'category') {
|
|
692
|
+
const Type = kind === 'float32' ? Float32Array : Float64Array;
|
|
693
|
+
const f = values instanceof Type ? values : Type.from(values);
|
|
694
|
+
buffers.push(new Uint8Array(0), asBytes(f));
|
|
695
|
+
return;
|
|
696
|
+
}
|
|
697
|
+
const map = this.labels[k], fresh = [];
|
|
698
|
+
let codes;
|
|
699
|
+
const code = (label) => {
|
|
700
|
+
let c = map.get(label);
|
|
701
|
+
if (c === undefined) { c = map.size; map.set(label, c); fresh.push(String(label)); }
|
|
702
|
+
return c;
|
|
703
|
+
};
|
|
704
|
+
if (values.codes) {
|
|
705
|
+
const lut = values.labels.map(code), src = values.codes;
|
|
706
|
+
if ((src instanceof Uint32Array || src instanceof Int32Array) && lut.every((c, j) => c === j)) codes = src; // same numbering: no copy
|
|
707
|
+
else { codes = new Int32Array(len); for (let i = 0; i < len; i++) codes[i] = lut[src[i]]; }
|
|
708
|
+
} else {
|
|
709
|
+
codes = new Int32Array(len);
|
|
710
|
+
for (let i = 0; i < len; i++) codes[i] = code(values[i]);
|
|
711
|
+
}
|
|
712
|
+
buffers.push(new Uint8Array(0), asBytes(codes));
|
|
713
|
+
if (fresh.length || !this.sent[k]) out.push(this.#dictionary(k, fresh));
|
|
714
|
+
});
|
|
715
|
+
const rb = recordBatch(Math.max(rows, 0), nodes, buffers);
|
|
716
|
+
out.push({ ...message(MESSAGE.batch, rb.header, rb.buffers, rb.bodyLength), dictionary: false });
|
|
717
|
+
return out;
|
|
718
|
+
}
|
|
719
|
+
|
|
720
|
+
#dictionary(k, labels) {
|
|
721
|
+
const parts = labels.map((s) => enc.encode(s));
|
|
722
|
+
const offsets = new Int32Array(labels.length + 1);
|
|
723
|
+
parts.forEach((p, i) => { offsets[i + 1] = offsets[i] + p.length; });
|
|
724
|
+
const rb = recordBatch(labels.length, [[labels.length, 0]], [new Uint8Array(0), asBytes(offsets), concatBytes(parts)]);
|
|
725
|
+
const header = tbl(['i64', k], ['off', rb.header], ['bool', this.sent[k]]);
|
|
726
|
+
this.sent[k] = true;
|
|
727
|
+
return { ...message(MESSAGE.dictionary, header, rb.buffers, rb.bodyLength), dictionary: true };
|
|
728
|
+
}
|
|
729
|
+
}
|
|
730
|
+
|
|
731
|
+
function concatBytes(list) {
|
|
732
|
+
const out = new Uint8Array(list.reduce((a, b) => a + b.length, 0));
|
|
733
|
+
let at = 0;
|
|
734
|
+
for (const b of list) { out.set(b, at); at += b.length; }
|
|
735
|
+
return out;
|
|
736
|
+
}
|
|
737
|
+
|
|
738
|
+
const isCoded = (v) => v != null && !ArrayBuffer.isView(v) && !Array.isArray(v) && v.codes != null && Array.isArray(v.labels);
|
|
739
|
+
|
|
740
|
+
// Write columns as Arrow IPC: the stream format, or with file: true the file
|
|
741
|
+
// format (.arrow / .feather, what pandas.read_feather and pyarrow open).
|
|
742
|
+
// columns: { name: Float64Array | Float32Array | number[] | string[] | { codes, labels } }.
|
|
743
|
+
// dates: names of columns holding dates (ms), written as Timestamp(ms, UTC).
|
|
744
|
+
export function writeArrow(columns, { file = false, dates = [] } = {}) {
|
|
745
|
+
const kinds = {};
|
|
746
|
+
for (const [name, v] of Object.entries(columns)) {
|
|
747
|
+
kinds[name] = dates.includes(name) ? 'date' : isCoded(v) || (Array.isArray(v) && typeof v[0] === 'string') ? 'category' : v instanceof Float32Array ? 'float32' : 'number';
|
|
748
|
+
}
|
|
749
|
+
const e = new ArrowEncoder(kinds);
|
|
750
|
+
const schema = e.schema(), parts = e.messages(columns);
|
|
751
|
+
if (!file) return concatBytes([schema, ...parts.map((p) => p.bytes), EOS]);
|
|
752
|
+
const head = new Uint8Array(8);
|
|
753
|
+
head.set(MAGIC);
|
|
754
|
+
let at = head.length + schema.length;
|
|
755
|
+
const dicts = [], batches = [];
|
|
756
|
+
for (const p of parts) {
|
|
757
|
+
(p.dictionary ? dicts : batches).push([at, p.metaLength, p.bodyLength]);
|
|
758
|
+
at += p.bytes.length;
|
|
759
|
+
}
|
|
760
|
+
const blocks = (list) => { // struct Block { offset: long; metaDataLength: int; bodyLength: long }
|
|
761
|
+
const dv = new DataView(new ArrayBuffer(24 * list.length));
|
|
762
|
+
list.forEach(([o, m, b], i) => { dv.setBigInt64(24 * i, BigInt(o), true); dv.setInt32(24 * i + 8, m, true); dv.setBigInt64(24 * i + 16, BigInt(b), true); });
|
|
763
|
+
return { structs: new Uint8Array(dv.buffer), count: list.length };
|
|
764
|
+
};
|
|
765
|
+
const footer = new FlatWriter().finish(tbl(['i16', V5], ['off', schemaTable(e.columns)], ['off', blocks(dicts)], ['off', blocks(batches)]));
|
|
766
|
+
const tail = new Uint8Array(10);
|
|
767
|
+
new DataView(tail.buffer).setInt32(0, footer.length, true);
|
|
768
|
+
tail.set(MAGIC, 4);
|
|
769
|
+
return concatBytes([head, schema, ...parts.map((p) => p.bytes), EOS, footer, tail]);
|
|
770
|
+
}
|
|
771
|
+
|
|
772
|
+
// ---- typed messages (Flight SQL) ----------------------------------------------------
|
|
773
|
+
// Arrow IPC messages for given column types, one message at a time (a Flight
|
|
774
|
+
// stream sends each as FlightData: its flatbuffer header and its body).
|
|
775
|
+
// fields: [{ name, type, nullable = true }], type one of
|
|
776
|
+
// utf8, binary, int8, int16, int32, int64, bool, float64,
|
|
777
|
+
// timestamp (ms, UTC), date32, dict (utf8 labels with int32 codes; one dictionary)
|
|
778
|
+
// Values: arrays; null / undefined (and NaN for numbers) are nulls. A dict
|
|
779
|
+
// column is { codes, labels } (label '' is null) or an array of strings.
|
|
780
|
+
const INT_BITS = { int8: 8, int16: 16, int32: 32, int64: 64 };
|
|
781
|
+
function typedField(f, k) {
|
|
782
|
+
const nullable = ['bool', f.nullable !== false];
|
|
783
|
+
if (f.type === 'dict') return tbl(['off', f.name], nullable, ['u8', T.Utf8], ['off', tbl()], ['off', tbl(['i64', k], ['off', tbl(['i32', 32], ['bool', 1])], ['bool', 0])], ['off', { vec: [] }]);
|
|
784
|
+
const [typeId, type] = INT_BITS[f.type] ? [T.Int, tbl(['i32', INT_BITS[f.type]], ['bool', 1])]
|
|
785
|
+
: f.type === 'float64' ? [T.FloatingPoint, tbl(['i16', 2])]
|
|
786
|
+
: f.type === 'utf8' ? [T.Utf8, tbl()] : f.type === 'binary' ? [T.Binary, tbl()] : f.type === 'bool' ? [T.Bool, tbl()]
|
|
787
|
+
: f.type === 'timestamp' ? [T.Timestamp, tbl(['i16', 1], ['off', 'UTC'])] : f.type === 'date32' ? [T.Date, tbl(['i16', 0])]
|
|
788
|
+
: (() => { throw new Error(`vectro: Arrow type ${f.type} is not supported here`); })();
|
|
789
|
+
return tbl(['off', f.name], nullable, ['u8', typeId], ['off', type], null, ['off', { vec: [] }]);
|
|
790
|
+
}
|
|
791
|
+
// The schema message ({ bytes, metaLength, bodyLength }).
|
|
792
|
+
export function typedSchema(fields) { return message(MESSAGE.schema, tbl(['i16', 0], ['off', { vec: fields.map(typedField) }])); }
|
|
793
|
+
|
|
794
|
+
const missingValue = (v) => v === null || v === undefined || (typeof v === 'number' && v !== v);
|
|
795
|
+
function validity(n, missing) { // a bitmap, or no buffer when nothing is missing
|
|
796
|
+
let nulls = 0;
|
|
797
|
+
const bits = new Uint8Array((n + 7) >> 3);
|
|
798
|
+
for (let i = 0; i < n; i++) { if (missing(i)) nulls++; else bits[i >> 3] |= 1 << (i & 7); }
|
|
799
|
+
return [nulls, nulls ? bits : new Uint8Array(0)];
|
|
800
|
+
}
|
|
801
|
+
function textBuffers(n, get) { // offsets and bytes of n strings (null: empty)
|
|
802
|
+
const parts = new Array(n), offsets = new Int32Array(n + 1);
|
|
803
|
+
for (let i = 0; i < n; i++) { const s = get(i); const b = s === null || s === undefined ? new Uint8Array(0) : typeof s === 'string' ? enc.encode(s) : s; parts[i] = b; offsets[i + 1] = offsets[i] + b.length; }
|
|
804
|
+
return [asBytes(offsets), concatBytes(parts)];
|
|
805
|
+
}
|
|
806
|
+
// The dictionary message of dict column k ({ bytes, metaLength, bodyLength }).
|
|
807
|
+
export function typedDictionary(k, labels) {
|
|
808
|
+
const [offs, data] = textBuffers(labels.length, (i) => labels[i]);
|
|
809
|
+
const rb = recordBatch(labels.length, [[labels.length, 0]], [new Uint8Array(0), offs, data]);
|
|
810
|
+
return message(MESSAGE.dictionary, tbl(['i64', k], ['off', rb.header], ['bool', 0]), rb.buffers, rb.bodyLength);
|
|
811
|
+
}
|
|
812
|
+
// (number columns in tight loops: no call per value)
|
|
813
|
+
function numberValidity(arr, from, n) {
|
|
814
|
+
let nulls = 0;
|
|
815
|
+
const bits = new Uint8Array((n + 7) >> 3);
|
|
816
|
+
for (let i = 0; i < n; i++) { const v = arr[from + i]; if (v === v) bits[i >> 3] |= 1 << (i & 7); else nulls++; }
|
|
817
|
+
return [nulls, nulls ? bits : new Uint8Array(0)];
|
|
818
|
+
}
|
|
819
|
+
function int64Words(arr, from, n) {
|
|
820
|
+
const w = new Int32Array(2 * n);
|
|
821
|
+
for (let i = 0; i < n; i++) { const v = arr[from + i]; if (v === v) { const x = Math.round(v), hi = Math.floor(x / 4294967296); w[2 * i] = x - hi * 4294967296; w[2 * i + 1] = hi; } }
|
|
822
|
+
return w;
|
|
823
|
+
}
|
|
824
|
+
function codedText(col, from, n) { // utf8 from { codes, labels }: label bytes made once
|
|
825
|
+
const bytes = col.labels.map((l) => enc.encode(l ?? '')), codes = col.codes, offsets = new Int32Array(n + 1);
|
|
826
|
+
const empty = Uint8Array.from(col.labels, (l) => (l === undefined || l === '' ? 1 : 0));
|
|
827
|
+
const lens = Int32Array.from(bytes, (b) => b.length);
|
|
828
|
+
const bits = new Uint8Array((n + 7) >> 3);
|
|
829
|
+
const nulls = textOffsets(codes, from, n, lens, empty, offsets, bits);
|
|
830
|
+
// all labels' bytes in one array, each row's copied byte by byte (short labels: faster than set())
|
|
831
|
+
const starts = new Int32Array(bytes.length + 1);
|
|
832
|
+
bytes.forEach((b, j) => { starts[j + 1] = starts[j] + b.length; });
|
|
833
|
+
const all = new Uint8Array(starts[bytes.length]);
|
|
834
|
+
bytes.forEach((b, j) => all.set(b, starts[j]));
|
|
835
|
+
const data = new Uint8Array(offsets[n]);
|
|
836
|
+
copyLabels(codes, from, n, starts, all, data, bytes.length);
|
|
837
|
+
return [nulls, nulls ? bits : new Uint8Array(0), asBytes(offsets), data];
|
|
838
|
+
}
|
|
839
|
+
|
|
840
|
+
function textOffsets(codes, from, n, lens, empty, offsets, bits) {
|
|
841
|
+
let nulls = 0, o = 0;
|
|
842
|
+
const L = lens.length;
|
|
843
|
+
for (let i = 0; i < n; i++) {
|
|
844
|
+
const c = codes[from + i];
|
|
845
|
+
if (c >= L || empty[c]) nulls++; else { bits[i >> 3] |= 1 << (i & 7); o += lens[c]; }
|
|
846
|
+
offsets[i + 1] = o;
|
|
847
|
+
}
|
|
848
|
+
return nulls;
|
|
849
|
+
}
|
|
850
|
+
function copyLabels(codes, from, n, starts, all, data, L) {
|
|
851
|
+
let o = 0;
|
|
852
|
+
for (let i = 0; i < n; i++) {
|
|
853
|
+
const c = codes[from + i];
|
|
854
|
+
if (c >= L) continue;
|
|
855
|
+
for (let j = starts[c], e = starts[c + 1]; j < e; j++) data[o++] = all[j];
|
|
856
|
+
}
|
|
857
|
+
}
|
|
858
|
+
|
|
859
|
+
// A record batch message of rows [from, to) of columns (arrays as above).
|
|
860
|
+
export function typedBatch(fields, columns, from, to) {
|
|
861
|
+
const n = to - from, nodes = [], buffers = [];
|
|
862
|
+
fields.forEach((f, k) => {
|
|
863
|
+
const col = columns[k], at = (i) => col[from + i];
|
|
864
|
+
if (col instanceof Float64Array && (f.type === 'float64' || f.type === 'int64' || f.type === 'timestamp')) {
|
|
865
|
+
const [nulls, valid] = numberValidity(col, from, n);
|
|
866
|
+
nodes.push([n, nulls]);
|
|
867
|
+
buffers.push(valid, f.type === 'float64' ? asBytes(col.subarray(from, to)) : asBytes(int64Words(col, from, n)));
|
|
868
|
+
return;
|
|
869
|
+
}
|
|
870
|
+
if (f.type === 'utf8' && col && col.codes) {
|
|
871
|
+
const [nulls, valid, offsets, data] = codedText(col, from, n);
|
|
872
|
+
nodes.push([n, nulls]);
|
|
873
|
+
buffers.push(valid, offsets, data);
|
|
874
|
+
return;
|
|
875
|
+
}
|
|
876
|
+
if (f.type === 'dict') {
|
|
877
|
+
const coded = col && col.codes, labels = coded ? col.labels : null;
|
|
878
|
+
const code = coded ? (i) => col.codes[from + i] : null;
|
|
879
|
+
const [nulls, valid] = validity(n, coded ? (i) => { const l = labels[code(i)]; return l === undefined || l === ''; } : (i) => missingValue(at(i)) || at(i) === '');
|
|
880
|
+
const codes = new Int32Array(n);
|
|
881
|
+
if (coded) for (let i = 0; i < n; i++) codes[i] = code(i);
|
|
882
|
+
else throw new Error('vectro: a dict column needs { codes, labels }');
|
|
883
|
+
nodes.push([n, nulls]); buffers.push(valid, asBytes(codes));
|
|
884
|
+
return;
|
|
885
|
+
}
|
|
886
|
+
const [nulls, valid] = validity(n, (i) => missingValue(at(i)));
|
|
887
|
+
nodes.push([n, nulls]);
|
|
888
|
+
if (f.type === 'utf8' || f.type === 'binary') { const [o, d] = textBuffers(n, at); buffers.push(valid, o, d); return; }
|
|
889
|
+
if (f.type === 'bool') { const bits = new Uint8Array((n + 7) >> 3); for (let i = 0; i < n; i++) if (!missingValue(at(i)) && at(i)) bits[i >> 3] |= 1 << (i & 7); buffers.push(valid, bits); return; }
|
|
890
|
+
if (f.type === 'float64') { const v = new Float64Array(n); for (let i = 0; i < n; i++) v[i] = missingValue(at(i)) ? 0 : at(i); buffers.push(valid, asBytes(v)); return; }
|
|
891
|
+
if (f.type === 'int64' || f.type === 'timestamp') {
|
|
892
|
+
const w = new Int32Array(2 * n);
|
|
893
|
+
for (let i = 0; i < n; i++) { const x = at(i); if (missingValue(x)) continue; const v = Math.round(typeof x === 'bigint' ? Number(x) : x), hi = Math.floor(v / 4294967296); w[2 * i] = v - hi * 4294967296; w[2 * i + 1] = hi; }
|
|
894
|
+
buffers.push(valid, asBytes(w));
|
|
895
|
+
return;
|
|
896
|
+
}
|
|
897
|
+
const Typed = f.type === 'int8' ? Int8Array : f.type === 'int16' ? Int16Array : Int32Array; // int8 / 16 / 32, date32 (days)
|
|
898
|
+
const v = new Typed(n);
|
|
899
|
+
for (let i = 0; i < n; i++) { const x = at(i); if (!missingValue(x)) v[i] = f.type === 'date32' ? Math.floor(x / 86_400_000) : x; }
|
|
900
|
+
buffers.push(valid, asBytes(v));
|
|
901
|
+
});
|
|
902
|
+
const rb = recordBatch(n, nodes, buffers);
|
|
903
|
+
return message(MESSAGE.batch, rb.header, rb.buffers, rb.bodyLength);
|
|
904
|
+
}
|
|
905
|
+
export { EOS as ARROW_EOS };
|