Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : // Arrow vectors assembled from their buffers.
       4              : // The Arrow builders (vectorFromArray, tableFromArrays of plain arrays, makeBuilder) compile their null check with
       5              : // `new Function`, which the Content-Security-Policy of the app forbids.
       6              : 
       7              : /// The validity bitmap of values, null and undefined are null
       8            1 : function validity(values: ReadonlyArray<unknown>): { nullCount: number; nullBitmap: Uint8Array | undefined } {
       9            1 :     const nullBitmap = new Uint8Array((values.length + 7) >> 3);
      10            1 :     let nullCount = 0;
      11            1 :     values.forEach((value, i) => {
      12            1 :         if (value == null) {
      13            1 :             ++nullCount;
      14              :         } else {
      15            1 :             nullBitmap[i >> 3] |= 1 << (i & 7);
      16              :         }
      17            1 :     });
      18            1 :     return { nullCount, nullBitmap: nullCount > 0 ? nullBitmap : undefined };
      19            1 : }
      20              : 
      21              : /// The 32-bit value offsets and the concatenated bytes of variable-width values
      22            1 : function variableWidthBuffers(values: ReadonlyArray<Uint8Array | null | undefined>): { valueOffsets: Int32Array; data: Uint8Array } {
      23            1 :     const valueOffsets = new Int32Array(values.length + 1);
      24            1 :     let end = 0;
      25            1 :     values.forEach((bytes, i) => {
      26            1 :         end += bytes?.length ?? 0;
      27            1 :         if (end > 0x7fffffff) {
      28            0 :             throw new RangeError(`${end} bytes exceed the 32-bit offsets of an Arrow vector`);
      29              :         }
      30            1 :         valueOffsets[i + 1] = end;
      31            1 :     });
      32            1 :     const data = new Uint8Array(end);
      33            1 :     values.forEach((bytes, i) => bytes && data.set(bytes, valueOffsets[i]));
      34            1 :     return { valueOffsets, data };
      35            1 : }
      36              : 
      37              : /// A Utf8 vector of strings
      38            1 : export function utf8Vector(values: ReadonlyArray<string | null | undefined>): arrow.Vector<arrow.Utf8> {
      39            1 :     const encoder = new TextEncoder();
      40            1 :     const buffers = variableWidthBuffers(values.map(value => value == null ? null : encoder.encode(value)));
      41            1 :     return arrow.makeVector(arrow.makeData({ type: new arrow.Utf8(), length: values.length, ...validity(values), ...buffers }));
      42            1 : }
      43              : 
      44              : /// A Binary vector of byte arrays
      45            1 : export function binaryVector(values: ReadonlyArray<Uint8Array | null | undefined>): arrow.Vector<arrow.Binary> {
      46            1 :     const buffers = variableWidthBuffers(values);
      47            1 :     return arrow.makeVector(arrow.makeData({ type: new arrow.Binary(), length: values.length, ...validity(values), ...buffers }));
      48            1 : }
      49              : 
      50              : /// A vector of numbers in a single or double precision float type
      51            1 : export function floatVector<T extends arrow.Float32 | arrow.Float64>(type: T, values: ReadonlyArray<number | null | undefined>): arrow.Vector<T> {
      52            1 :     const data = new type.ArrayType(values.length);
      53            1 :     values.forEach((value, i) => (data[i] = value ?? 0));
      54            1 :     return arrow.makeVector(arrow.makeData({ type, length: values.length, ...validity(values), data }));
      55            1 : }
      56              : 
      57              : /// A vector of integers, 64-bit types take bigints
      58            1 : export function intVector<T extends arrow.Int>(type: T, values: ReadonlyArray<number | bigint | null | undefined>): arrow.Vector<T> {
      59            1 :     const data = new type.ArrayType(values.length);
      60            1 :     const zero = type.bitWidth === 64 ? 0n : 0;
      61            1 :     values.forEach((value, i) => ((data as unknown as Array<number | bigint>)[i] = value ?? zero));
      62            1 :     return arrow.makeVector(arrow.makeData({ type, length: values.length, ...validity(values), data }));
      63            1 : }
      64              : 
      65              : /// A Bool vector of booleans
      66            1 : export function boolVector(values: ReadonlyArray<boolean | null | undefined>): arrow.Vector<arrow.Bool> {
      67            1 :     const data = new Uint8Array((values.length + 7) >> 3);
      68            1 :     values.forEach((value, i) => {
      69            1 :         if (value) data[i >> 3] |= 1 << (i & 7);
      70            1 :     });
      71            1 :     return arrow.makeVector(arrow.makeData({ type: new arrow.Bool(), length: values.length, ...validity(values), data }));
      72            1 : }