Line data Source code
1 : import * as arrow from 'apache-arrow';
2 : import { Logger } from '../../platform/logger/logger.js';
3 :
4 : /// Create a validity bitmap from Uint8Array storing null indicators
5 : /// Returns [bitmap, nullCount]
6 1 : function createValidityBitmap(n: number, isNull: Uint8Array): [Uint8Array, number] {
7 1 : const validityBytes = Math.max(1, ((n + 63) & ~63) >> 3);
8 1 : const bitmap = new Uint8Array(validityBytes).fill(255); // all valid initially
9 1 : let nullCount = 0;
10 :
11 1 : for (let i = 0; i < n; i++) {
12 1 : if (isNull[i]) {
13 1 : const byte = i >> 3;
14 1 : const bit = i & 7;
15 1 : bitmap[byte] &= ~(1 << bit); // clear bit = null
16 1 : nullCount++;
17 1 : }
18 : }
19 1 : return [bitmap, nullCount];
20 1 : }
21 :
22 : /// Create Arrow Data for boolean column
23 1 : function createBoolData(type: arrow.Bool, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.Bool> {
24 1 : const n = values.length;
25 1 : tmpIsNull.fill(0);
26 1 : const byteLen = Math.max(1, (n + 7) >> 3);
27 1 : const buffer = new Uint8Array(byteLen);
28 :
29 1 : for (let i = 0; i < n; i++) {
30 1 : const v = values[i];
31 1 : if (v == null) {
32 1 : tmpIsNull[i] = 1;
33 1 : } else if (Boolean(v)) {
34 1 : buffer[i >> 3] |= (1 << (i & 7));
35 : }
36 1 : }
37 :
38 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
39 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
40 1 : }
41 :
42 : /// Create Arrow Data for Int types
43 1 : function createIntData<T extends arrow.Int>(type: T, values: any[], tmpIsNull: Uint8Array): arrow.Data<T> {
44 1 : const n = values.length;
45 1 : tmpIsNull.fill(0);
46 1 : const buffer = new type.ArrayType(n);
47 :
48 1 : for (let i = 0; i < n; i++) {
49 1 : const v = values[i];
50 1 : if (v == null || v === '') {
51 1 : tmpIsNull[i] = 1;
52 1 : } else {
53 1 : const num = Number(v);
54 1 : if (Number.isNaN(num)) {
55 1 : tmpIsNull[i] = 1;
56 : } else {
57 1 : buffer[i] = num;
58 : }
59 1 : }
60 1 : }
61 :
62 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
63 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
64 1 : }
65 :
66 : /// Create Arrow Data for Float types
67 1 : function createFloatData<T extends arrow.Float>(type: T, values: any[], tmpIsNull: Uint8Array): arrow.Data<T> {
68 1 : const n = values.length;
69 1 : tmpIsNull.fill(0);
70 1 : const buffer = new type.ArrayType(n);
71 :
72 1 : for (let i = 0; i < n; i++) {
73 1 : const v = values[i];
74 1 : if (v == null || v === '') {
75 1 : tmpIsNull[i] = 1;
76 1 : } else {
77 1 : const num = Number(v);
78 1 : if (Number.isNaN(num)) {
79 1 : tmpIsNull[i] = 1;
80 : } else {
81 1 : buffer[i] = num;
82 : }
83 1 : }
84 1 : }
85 :
86 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
87 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
88 1 : }
89 :
90 : /// Create Arrow Data for Int64/BigInt types
91 1 : function createInt64Data(type: arrow.Int64, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.Int64> {
92 1 : const n = values.length;
93 1 : tmpIsNull.fill(0);
94 1 : const buffer = new type.ArrayType(n);
95 :
96 1 : for (let i = 0; i < n; i++) {
97 1 : const v = values[i];
98 1 : if (v == null || v === '') {
99 1 : tmpIsNull[i] = 1;
100 : } else {
101 1 : try {
102 1 : buffer[i] = BigInt(v);
103 1 : } catch {
104 1 : tmpIsNull[i] = 1;
105 1 : }
106 : }
107 1 : }
108 :
109 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
110 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
111 1 : }
112 :
113 : /// Create Arrow Data for Decimal128
114 1 : function createDecimalData(type: arrow.Decimal, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.Decimal> {
115 1 : const n = values.length;
116 1 : const stride = type.bitWidth / 32;
117 1 : const buffer = new Uint32Array(n * stride);
118 1 : tmpIsNull.fill(0);
119 :
120 1 : for (let i = 0; i < n; i++) {
121 1 : const value = values[i];
122 1 : if (value == null || value === '') {
123 0 : tmpIsNull[i] = 1;
124 0 : continue;
125 0 : }
126 :
127 1 : try {
128 1 : const text = String(value).trim();
129 1 : const match = /^([+-]?)(\d*)(?:\.(\d*))?$/.exec(text);
130 1 : if (!match || (!match[2] && !match[3])) {
131 1 : tmpIsNull[i] = 1;
132 1 : continue;
133 1 : }
134 :
135 1 : const fraction = match[3] ?? '';
136 1 : if (fraction.length > type.scale) {
137 1 : tmpIsNull[i] = 1;
138 1 : continue;
139 1 : }
140 1 : const digits = `${match[2] || '0'}${fraction.padEnd(type.scale, '0')}`.replace(/^0+(?=\d)/, '');
141 1 : let unscaled = BigInt(digits || '0');
142 1 : if (match[1] === '-') unscaled = -unscaled;
143 :
144 1 : const maxValue = (1n << BigInt(type.bitWidth - 1)) - 1n;
145 1 : const minValue = -(1n << BigInt(type.bitWidth - 1));
146 1 : if (unscaled < minValue || unscaled > maxValue) {
147 1 : tmpIsNull[i] = 1;
148 1 : continue;
149 1 : }
150 :
151 1 : if (unscaled < 0) unscaled += 1n << BigInt(type.bitWidth);
152 1 : for (let word = 0; word < stride; word++) {
153 1 : buffer[i * stride + word] = Number(unscaled & 0xffffffffn);
154 1 : unscaled >>= 32n;
155 1 : }
156 1 : } catch {
157 0 : tmpIsNull[i] = 1;
158 0 : }
159 1 : }
160 :
161 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
162 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
163 1 : }
164 :
165 : /// Create Arrow Data for DateDay (days since epoch as Int32)
166 1 : function createDateDayData(type: arrow.DateDay, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.DateDay> {
167 1 : const n = values.length;
168 1 : tmpIsNull.fill(0);
169 1 : const buffer = new Int32Array(n);
170 :
171 1 : for (let i = 0; i < n; i++) {
172 1 : const v = values[i];
173 1 : if (v == null || v === '') {
174 1 : tmpIsNull[i] = 1;
175 1 : } else {
176 1 : const date = new Date(v);
177 1 : const ms = date.getTime();
178 1 : if (Number.isNaN(ms)) {
179 1 : tmpIsNull[i] = 1;
180 : } else {
181 1 : buffer[i] = Math.floor(ms / (24 * 60 * 60 * 1000));
182 : }
183 1 : }
184 1 : }
185 :
186 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
187 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
188 1 : }
189 :
190 : /// Create Arrow Data for DateMillisecond (ms since epoch as BigInt64)
191 1 : function createDateMillisecondData(type: arrow.DateMillisecond, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.DateMillisecond> {
192 1 : const n = values.length;
193 1 : tmpIsNull.fill(0);
194 1 : const buffer = new BigInt64Array(n);
195 :
196 1 : for (let i = 0; i < n; i++) {
197 1 : const v = values[i];
198 1 : if (v == null || v === '') {
199 1 : tmpIsNull[i] = 1;
200 1 : } else {
201 1 : const date = new Date(v);
202 1 : const ms = date.getTime();
203 1 : if (Number.isNaN(ms)) {
204 0 : tmpIsNull[i] = 1;
205 : } else {
206 1 : buffer[i] = BigInt(ms);
207 : }
208 1 : }
209 1 : }
210 :
211 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
212 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
213 1 : }
214 :
215 : /// Create Arrow Data for TimeMillisecond (ms since midnight as Int32)
216 1 : function createTimeMillisecondData(type: arrow.TimeMillisecond, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.TimeMillisecond> {
217 1 : const n = values.length;
218 1 : tmpIsNull.fill(0);
219 1 : const buffer = new Int32Array(n);
220 :
221 1 : for (let i = 0; i < n; i++) {
222 1 : const v = values[i];
223 1 : if (v == null) {
224 1 : tmpIsNull[i] = 1;
225 1 : } else {
226 : // Parse time string "HH:MM:SS.sss"
227 1 : const parts = String(v).split(':');
228 1 : const hours = parseInt(parts[0], 10) || 0;
229 1 : const minutes = parseInt(parts[1], 10) || 0;
230 1 : const secondsParts = (parts[2] || '0').split('.');
231 1 : const seconds = parseInt(secondsParts[0], 10) || 0;
232 1 : const millis = parseInt((secondsParts[1] || '0').padEnd(3, '0').slice(0, 3), 10);
233 1 : buffer[i] = ((hours * 60 + minutes) * 60 + seconds) * 1000 + millis;
234 1 : }
235 1 : }
236 :
237 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
238 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
239 1 : }
240 :
241 : /// A Trino timestamp: `2024-03-01 12:30:00.123`, with time zone followed by ` UTC`, ` +01:00` or ` Europe/Berlin`
242 1 : const TRINO_TIMESTAMP_PATTERN = /^(-?\d+)-(\d{2})-(\d{2})[ T](\d{2}):(\d{2}):(\d{2})(?:\.(\d+))?(?: ?(\S+))?$/;
243 : /// A zone offset like `+01:00`
244 1 : const ZONE_OFFSET_PATTERN = /^([+-])(\d{2}):?(\d{2})$/;
245 : /// The zones that are UTC
246 1 : const UTC_ZONES = new Set(['UTC', 'Z', 'GMT', 'Etc/UTC', 'Etc/GMT', 'UCT', 'Zulu', 'Universal']);
247 :
248 : /// The milliseconds since epoch of a wall clock time in UTC, Date.UTC would move the years 0 to 99 to the 1900s
249 1 : function utcMs(year: number, month: number, day: number, hours: number, minutes: number, seconds: number, ms: number): number {
250 1 : const date = new Date(0);
251 1 : date.setUTCFullYear(year, month - 1, day);
252 1 : date.setUTCHours(hours, minutes, seconds, ms);
253 1 : return date.getTime();
254 1 : }
255 :
256 : /// The formatters of the wall clock in a named zone
257 1 : const zoneFormatters = new Map<string, Intl.DateTimeFormat>();
258 : /// The offsets of named zones by zone and minute, cleared when it grows too large
259 1 : const zoneOffsets = new Map<string, number>();
260 1 : const ZONE_OFFSET_CACHE_SIZE = 100_000;
261 :
262 : /// The offset of a named zone at an instant: the wall clock in the zone minus the wall clock in UTC
263 1 : function namedZoneOffsetMs(zone: string, instantMs: number): number {
264 1 : const minute = Math.floor(instantMs / 60_000);
265 1 : const key = `${zone}|${minute}`;
266 1 : const cached = zoneOffsets.get(key);
267 1 : if (cached !== undefined) return cached;
268 :
269 1 : let formatter = zoneFormatters.get(zone);
270 1 : if (formatter === undefined) {
271 : // Throws a RangeError for unknown zones
272 1 : formatter = new Intl.DateTimeFormat('en-US', {
273 1 : timeZone: zone, hourCycle: 'h23', era: 'short',
274 1 : year: 'numeric', month: 'numeric', day: 'numeric', hour: 'numeric', minute: 'numeric',
275 1 : });
276 1 : zoneFormatters.set(zone, formatter);
277 1 : }
278 1 : const parts: Record<string, string> = {};
279 1 : for (const part of formatter.formatToParts(minute * 60_000)) parts[part.type] = part.value;
280 1 : const year = parts.era === 'BC' ? 1 - Number(parts.year) : Number(parts.year);
281 1 : const wall = utcMs(year, Number(parts.month), Number(parts.day), Number(parts.hour), Number(parts.minute), 0, 0);
282 1 : const offset = wall - minute * 60_000;
283 :
284 1 : if (zoneOffsets.size >= ZONE_OFFSET_CACHE_SIZE) zoneOffsets.clear();
285 1 : zoneOffsets.set(key, offset);
286 1 : return offset;
287 1 : }
288 :
289 : /// Parse a Trino timestamp to milliseconds since epoch, null if it is invalid.
290 : /// A timestamp without time zone is a wall clock time and stays as is, one with a time zone is moved to UTC.
291 1 : export function parseTrinoTimestampMs(text: string): number | null {
292 1 : const match = TRINO_TIMESTAMP_PATTERN.exec(text);
293 1 : if (match == null) return null;
294 1 : const fraction = (match[7] ?? '').padEnd(3, '0').slice(0, 3);
295 1 : const wall = utcMs(Number(match[1]), Number(match[2]), Number(match[3]), Number(match[4]), Number(match[5]), Number(match[6]), Number(fraction));
296 1 : const zone = match[8];
297 1 : if (zone == null || UTC_ZONES.has(zone)) {
298 1 : return wall;
299 : }
300 1 : const offset = ZONE_OFFSET_PATTERN.exec(zone);
301 1 : if (offset != null) {
302 1 : const offsetMs = (Number(offset[2]) * 60 + Number(offset[3])) * 60_000;
303 1 : return offset[1] === '-' ? wall + offsetMs : wall - offsetMs;
304 1 : }
305 1 : try {
306 : // Guess the instant with the offset at the wall clock, then correct it with the offset at the guess
307 1 : const guess = wall - namedZoneOffsetMs(zone, wall);
308 1 : return wall - namedZoneOffsetMs(zone, guess);
309 1 : } catch {
310 1 : return null;
311 1 : }
312 1 : }
313 :
314 : /// Create Arrow Data for Timestamp types (ms since epoch as BigInt64)
315 1 : function createTimestampData<T extends arrow.Timestamp>(type: T, values: any[], tmpIsNull: Uint8Array): arrow.Data<T> {
316 1 : const n = values.length;
317 1 : tmpIsNull.fill(0);
318 1 : const buffer = new BigInt64Array(n);
319 :
320 1 : for (let i = 0; i < n; i++) {
321 1 : const v = values[i];
322 1 : const ms = v == null || v === '' ? null : parseTrinoTimestampMs(String(v));
323 1 : if (ms == null || Number.isNaN(ms)) {
324 1 : tmpIsNull[i] = 1;
325 : } else {
326 1 : buffer[i] = BigInt(ms);
327 : }
328 1 : }
329 :
330 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
331 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, data: buffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
332 1 : }
333 :
334 : // Shared TextEncoder for string encoding
335 1 : const textEncoder = new TextEncoder();
336 :
337 : /// Create Arrow Data for Utf8 strings
338 1 : function createUtf8Data(type: arrow.Utf8, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.Utf8> {
339 1 : const n = values.length;
340 1 : tmpIsNull.fill(0);
341 1 : const encodedValues = new Array<Uint8Array | null>(n);
342 1 : let totalBytes = 0;
343 :
344 : // First pass: encode strings and track nulls
345 1 : for (let i = 0; i < n; i++) {
346 1 : const v = values[i];
347 1 : if (v == null) {
348 1 : tmpIsNull[i] = 1;
349 1 : encodedValues[i] = null;
350 1 : } else {
351 1 : const str = typeof v === 'object' ? JSON.stringify(v) : String(v);
352 1 : const encoded = textEncoder.encode(str);
353 1 : encodedValues[i] = encoded;
354 1 : totalBytes += encoded.length;
355 1 : }
356 1 : }
357 :
358 : // Second pass: build offsets and data buffer
359 1 : const offsets = new Int32Array(n + 1);
360 1 : const dataBuffer = new Uint8Array(totalBytes);
361 1 : let offset = 0;
362 :
363 1 : for (let i = 0; i < n; i++) {
364 1 : offsets[i] = offset;
365 1 : const encoded = encodedValues[i];
366 1 : if (encoded != null) {
367 1 : dataBuffer.set(encoded, offset);
368 1 : offset += encoded.length;
369 1 : }
370 1 : }
371 1 : offsets[n] = offset;
372 :
373 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
374 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, valueOffsets: offsets, data: dataBuffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
375 1 : }
376 :
377 : /// Create Arrow Data for Binary
378 1 : function createBinaryData(type: arrow.Binary, values: any[], tmpIsNull: Uint8Array): arrow.Data<arrow.Binary> {
379 1 : const n = values.length;
380 1 : tmpIsNull.fill(0);
381 1 : const binaryValues = new Array<Uint8Array | null>(n);
382 1 : let totalBytes = 0;
383 :
384 : // First pass: decode binary values and track nulls
385 1 : for (let i = 0; i < n; i++) {
386 1 : const v = values[i];
387 1 : if (v == null) {
388 0 : tmpIsNull[i] = 1;
389 0 : binaryValues[i] = null;
390 1 : } else if (v instanceof Uint8Array) {
391 1 : binaryValues[i] = v;
392 1 : totalBytes += v.length;
393 1 : } else if (typeof v === 'string') {
394 : // Trino returns varbinary as base64
395 1 : const bytes = Uint8Array.from(atob(v), c => c.charCodeAt(0));
396 1 : binaryValues[i] = bytes;
397 1 : totalBytes += bytes.length;
398 1 : } else {
399 1 : tmpIsNull[i] = 1;
400 1 : binaryValues[i] = null;
401 1 : }
402 1 : }
403 :
404 : // Second pass: build offsets and data buffer
405 1 : const offsets = new Int32Array(n + 1);
406 1 : const dataBuffer = new Uint8Array(totalBytes);
407 1 : let offset = 0;
408 :
409 1 : for (let i = 0; i < n; i++) {
410 1 : offsets[i] = offset;
411 1 : const bytes = binaryValues[i];
412 1 : if (bytes != null) {
413 1 : dataBuffer.set(bytes, offset);
414 1 : offset += bytes.length;
415 1 : }
416 1 : }
417 1 : offsets[n] = offset;
418 :
419 1 : const [validityBitmap, nullCount] = createValidityBitmap(n, tmpIsNull);
420 1 : return arrow.makeData({ type, offset: 0, length: n, nullCount, valueOffsets: offsets, data: dataBuffer, nullBitmap: nullCount > 0 ? validityBitmap : undefined });
421 1 : }
422 :
423 : /// Create Arrow Data for a column based on its type.
424 : /// Instances report the type family as typeId (Type.Int, Type.Date, Type.Time), never the flattened ids like
425 : /// Type.Int32 or Type.DateDay, so we dispatch on the family and then on the bit width or the unit.
426 1 : function createColumnData(field: arrow.Field, values: any[], tmpIsNull: Uint8Array): arrow.Data {
427 1 : const type = field.type;
428 1 : if (arrow.DataType.isBool(type)) {
429 1 : return createBoolData(type, values, tmpIsNull);
430 : }
431 1 : if (arrow.DataType.isInt(type)) {
432 1 : return type.bitWidth === 64
433 1 : ? createInt64Data(type as arrow.Int64, values, tmpIsNull)
434 1 : : createIntData(type, values, tmpIsNull);
435 : }
436 1 : if (arrow.DataType.isFloat(type)) {
437 1 : return createFloatData(type, values, tmpIsNull);
438 : }
439 1 : if (arrow.DataType.isDecimal(type)) {
440 1 : return createDecimalData(type, values, tmpIsNull);
441 : }
442 1 : if (arrow.DataType.isDate(type)) {
443 1 : return type.unit === arrow.DateUnit.DAY
444 1 : ? createDateDayData(type as arrow.DateDay, values, tmpIsNull)
445 1 : : createDateMillisecondData(type as arrow.DateMillisecond, values, tmpIsNull);
446 : }
447 1 : if (arrow.DataType.isTime(type) && type.unit === arrow.TimeUnit.MILLISECOND && type.bitWidth === 32) {
448 1 : return createTimeMillisecondData(type as arrow.TimeMillisecond, values, tmpIsNull);
449 : }
450 1 : if (arrow.DataType.isTimestamp(type)) {
451 1 : return createTimestampData(type, values, tmpIsNull);
452 : }
453 1 : if (arrow.DataType.isBinary(type)) {
454 1 : return createBinaryData(type, values, tmpIsNull);
455 : }
456 : // Utf8 and all other types (including complex types as JSON)
457 1 : return createUtf8Data(new arrow.Utf8(), values, tmpIsNull);
458 1 : }
459 :
460 : /// Translate any[] with a schema to an arrow batch
461 1 : export function translateAnyRowsToArrowBatch(schema: arrow.Schema, rows: any[][], _logger: Logger): arrow.RecordBatch {
462 1 : const numRows = rows.length;
463 1 : const numCols = schema.fields.length;
464 :
465 : // Shared validity buffer reused across all columns
466 1 : const tmpIsNull = new Uint8Array(numRows);
467 :
468 : // Build column data directly
469 1 : const columnData: arrow.Data[] = [];
470 :
471 :
472 1 : for (let col = 0; col < numCols; col++) {
473 1 : const field = schema.fields[col];
474 :
475 : // Collect column values
476 1 : const values: any[] = [];
477 1 : for (let row = 0; row < numRows; row++) {
478 1 : values.push(rows[row][col]);
479 : }
480 :
481 : // Create Data for this column
482 1 : const data = createColumnData(field, values, tmpIsNull);
483 1 : columnData.push(data);
484 1 : }
485 :
486 : // Create struct and batch
487 1 : const structData = arrow.makeData({ type: new arrow.Struct(schema.fields), children: columnData, nullCount: 0 });
488 1 : return new arrow.RecordBatch(schema, structData);
489 1 : }
|
