arquero-arrow
Version:
Arrow serialization support for Arquero.
647 lines (576 loc) • 19 kB
JavaScript
(function (global, factory) {
typeof exports === 'object' && typeof module !== 'undefined' ? factory(exports, require('apache-arrow')) :
typeof define === 'function' && define.amd ? define(['exports', 'apache-arrow'], factory) :
(global = typeof globalThis !== 'undefined' ? globalThis : global || self, factory((global.aq = global.aq || {}, global.aq.arrow = {}), global.Arrow));
}(this, (function (exports, apacheArrow) { 'use strict';
function ceil64Bytes(length, bpe = 1) {
return ((((length * bpe) + 63) & ~63) || 64) / bpe;
}
function array(Type, length, bpe = Type.BYTES_PER_ELEMENT) {
return new Type(ceil64Bytes(length, bpe));
}
function arrowData(d) {
return d instanceof apacheArrow.Data
? d
: apacheArrow.Data.new(d.type, 0, d.length, d.nulls, d.buffers, null, d.dict);
}
function arrowVector(data) {
return apacheArrow.Vector.new(arrowData(data));
}
const encoder = new TextEncoder();
function encode(data, idx, str) {
const bytes = encoder.encode(str);
data.set(bytes, idx);
return bytes.length;
}
function encodeInto(data, idx, str) {
return encoder.encodeInto(str, data.subarray(idx)).written;
}
const writeUtf8 = encoder.encodeInto ? encodeInto : encode;
function arrayBuilder(type, length) {
const data = array(type.ArrayType, length);
return {
set(value, index) { data[index] = value; },
data: () => ({ type, length, buffers: [null, data] })
};
}
function boolBuilder(type, length) {
const data = array(type.ArrayType, length / 8);
return {
set(value, index) {
if (value) data[index >> 3] |= (1 << (index % 8));
},
data: () => ({ type, length, buffers: [null, data] })
};
}
function dateDayBuilder(type, length) {
const data = array(type.ArrayType, length);
return {
set(value, index) { data[index] = (value / 86400000) | 0; },
data: () => ({ type, length, buffers: [null, data] })
};
}
function dateMillisBuilder(type, length) {
const data = array(type.ArrayType, length);
return {
set(value, index) {
const i = index << 1;
data[ i] = (value % 4294967296) | 0;
data[i+1] = (value / 4294967296) | 0;
},
data: () => ({ type, length, buffers: [null, data] })
};
}
function defaultBuilder(type) {
const b = apacheArrow.Builder.new({
type,
nullValues: [null, undefined],
highWaterMark: Infinity
});
return {
set(value, index) { b.set(index, value); },
data: () => b.finish().flush()
};
}
function utf8Builder(type, length, strlen) {
const offset = array(Int32Array, length + 1);
const buf = array(Uint8Array, 3 * strlen);
let idx = 0;
return {
set(value, index) {
idx += writeUtf8(buf, idx, value);
offset[index + 1] = idx;
},
data: () => {
// slice utf buffer if over-allocated
const dlen = ceil64Bytes(idx);
const data = buf.length > dlen ? buf.subarray(0, dlen) : buf;
return { type, length, buffers: [offset, data] };
}
};
}
function dictionaryBuilder(type, length) {
const values = [];
const data = array(type.indices.ArrayType, length);
const keys = Object.create(null);
let next = -1;
let strlen = 0;
return {
set(value, index) {
const v = String(value);
let k = keys[v];
if (k === undefined) {
strlen += v.length;
keys[v] = k = ++next;
values.push(v);
}
data[index] = k;
},
data: () => ({
type,
length,
buffers: [null, data],
dict: dictionary(type.dictionary, values, strlen)
})
};
}
function dictionary(type, values, strlen) {
const b = utf8Builder(type, values.length, strlen);
values.forEach(b.set);
return arrowVector(b.data());
}
function validBuilder(builder, length) {
const valid = array(Uint8Array, length / 8);
let nulls = 0;
return {
set(value, index) {
if (value == null) {
++nulls;
} else {
builder.set(value, index);
valid[index >> 3] |= (1 << (index % 8));
}
},
data: () => {
const d = builder.data();
if (nulls) {
d.nulls = nulls;
d.buffers[2] = valid;
}
return d;
}
};
}
function builder(type, nrows, nullable = true) {
let method;
switch (type.typeId) {
case apacheArrow.Type.Int:
method = type.bitWidth < 64 ? arrayBuilder : null;
break;
case apacheArrow.Type.Float:
method = type.precision > 0 ? arrayBuilder : null;
break;
case apacheArrow.Type.Dictionary:
// check sub-types against builder assumptions
// if check fails, fallback to default builder
method = (
type.dictionary.typeId === apacheArrow.Type.Utf8 &&
type.indices.typeId === apacheArrow.Type.Int &&
type.indices.bitWidth < 64
) ? dictionaryBuilder : null;
break;
case apacheArrow.Type.Bool:
method = boolBuilder;
break;
case apacheArrow.Type.Date:
method = type.unit ? dateMillisBuilder : dateDayBuilder;
break;
}
return method == null ? defaultBuilder(type)
: nullable ? validBuilder(method(type, nrows), nrows)
: method(type, nrows);
}
function dataFromArray(array, type) {
const length = array.length;
const size = ceil64Bytes(length, array.BYTES_PER_ELEMENT);
let data = array;
if (length !== size) {
data = new array.constructor(size);
data.set(array);
}
return arrowData({ type, length, buffers: [null, data] });
}
function dataFromScan(nrows, scan, column, type, nullable = true) {
const b = builder(type, nrows, nullable);
scan(column, b.set);
return arrowData(b.data());
}
var isArray = Array.isArray;
function isTypedArray(value) {
// all typed arrays should share the same method prototype
return value && value.map === Int8Array.prototype.map;
}
function isArrayType(value) {
return isArray(value) || isTypedArray(value);
}
function scanTable(table, limit, offset) {
const scanAll = offset === 0 && table.numRows() <= limit
&& !table.isFiltered() && !table.isOrdered();
return (column, visit) => {
let i = -1;
scanAll && isArrayType(column.data)
? column.data.forEach(visit)
: table.scan(
row => visit(column.get(row), ++i),
true, limit, offset
);
};
}
function error(message) {
throw Error(message);
}
function toString(v) {
return v === undefined ? v + ''
: typeof v === 'bigint' ? v + 'n'
: JSON.stringify(v);
}
function resolveType(type) {
if (type instanceof apacheArrow.DataType || type == null) {
return type;
}
switch (type) {
case apacheArrow.Type.Binary:
return new apacheArrow.Binary();
case apacheArrow.Type.Bool:
return new apacheArrow.Bool();
case apacheArrow.Type.DateDay:
return new apacheArrow.DateDay();
case apacheArrow.Type.DateMillisecond:
case apacheArrow.Type.Date:
return new apacheArrow.DateMillisecond();
case apacheArrow.Type.Dictionary:
return new apacheArrow.Dictionary(new apacheArrow.Utf8(), new apacheArrow.Int32());
case apacheArrow.Type.Float16:
return new apacheArrow.Float16();
case apacheArrow.Type.Float32:
return new apacheArrow.Float32();
case apacheArrow.Type.Float64:
case apacheArrow.Type.Float:
return new apacheArrow.Float64();
case apacheArrow.Type.Int8:
return new apacheArrow.Int8();
case apacheArrow.Type.Int16:
return new apacheArrow.Int16();
case apacheArrow.Type.Int32:
case apacheArrow.Type.Int:
return new apacheArrow.Int32();
case apacheArrow.Type.Int64:
return new apacheArrow.Int64();
case apacheArrow.Type.IntervalDayTime:
return new apacheArrow.IntervalDayTime();
case apacheArrow.Type.Interval:
case apacheArrow.Type.IntervalYearMonth:
return new apacheArrow.IntervalYearMonth();
case apacheArrow.Type.Null:
return new apacheArrow.Null();
case apacheArrow.Type.TimeMicrosecond:
return new apacheArrow.TimeMicrosecond();
case apacheArrow.Type.TimeMillisecond:
case apacheArrow.Type.Time:
return new apacheArrow.TimeMillisecond();
case apacheArrow.Type.TimeNanosecond:
return new apacheArrow.TimeNanosecond();
case apacheArrow.Type.TimeSecond:
return new apacheArrow.TimeSecond();
case apacheArrow.Type.Uint8:
return new apacheArrow.Uint8();
case apacheArrow.Type.Uint16:
return new apacheArrow.Uint16();
case apacheArrow.Type.Uint32:
return new apacheArrow.Uint32();
case apacheArrow.Type.Uint64:
return new apacheArrow.Uint64();
case apacheArrow.Type.Utf8:
return new apacheArrow.Utf8();
default:
error(
`Unsupported type code: ${toString(type)}. ` +
'Use a data type constructor instead?'
);
}
}
function isDate(value) {
return value instanceof Date;
}
function isExactUTCDate(d) {
return d.getUTCHours() === 0
&& d.getUTCMinutes() === 0
&& d.getUTCSeconds() === 0
&& d.getUTCMilliseconds() === 0;
}
function profile(scan, column) {
const p = profiler();
scan(column, p.add);
return p;
}
function profiler() {
const p = {
count: 0,
nulls: 0,
bools: 0,
nums: 0,
ints: 0,
bigints: 0,
min: Infinity,
max: -Infinity,
digits: 0,
dates: 0,
utcdays: 0,
strings: 0,
strlen: 0,
arrays: 0,
minlen: Infinity,
maxlen: 0,
structs: 0,
add(value) {
++p.count;
if (value == null) {
++p.nulls;
return;
}
const type = typeof value;
if (type === 'string') {
++p.strings;
} else if (type === 'number') {
++p.nums;
if (value < p.min) p.min = value;
if (value > p.max) p.max = value;
if (Number.isInteger(value)) ++p.ints;
} else if (type === 'boolean') {
++p.bools;
} else if (type === 'object') {
if (isDate(value)) {
++p.dates;
if (isExactUTCDate(value)) {
++p.utcdays;
}
} else if (isArrayType(value)) {
++p.arrays;
if (value.length < p.minlen) p.minlen = value.length;
if (value.length > p.maxlen) p.maxlen = value.length;
const ap = p.array_prof || (p.array_prof = profiler());
value.forEach(ap.add);
} else {
++p.structs;
const sp = p.struct_prof || (p.struct_prof = {});
for (const key in value) {
const fp = sp[key] || (sp[key] = profiler());
fp.add(value[key]);
}
}
} else if (type === 'bigint') {
++p.bigints;
if (value < p.min) p.min = value;
if (value > p.max) p.max = value;
}
},
type() {
return resolveType(infer(p));
}
};
return p;
}
function infer(p) {
const valid = p.count - p.nulls;
if (valid === 0) {
return apacheArrow.Type.Null;
}
else if (p.ints === valid) {
const v = Math.max(Math.abs(p.min) - 1, p.max);
return p.min < 0
? v >= 2 ** 31 ? apacheArrow.Type.Float64
: v < (1 << 7) ? apacheArrow.Type.Int8 : v < (1 << 15) ? apacheArrow.Type.Int16 : apacheArrow.Type.Int32
: v >= 2 ** 32 ? apacheArrow.Type.Float64
: v < (1 << 8) ? apacheArrow.Type.Uint8 : v < (1 << 16) ? apacheArrow.Type.Uint16 : apacheArrow.Type.Uint32;
}
else if (p.nums === valid) {
return apacheArrow.Type.Float64;
}
else if (p.bigints === valid) {
const v = -p.min > p.max ? -p.min - 1n : p.max;
return p.min < 0
? v < 2 ** 63 ? apacheArrow.Type.Int64
: error(`BigInt exceeds 64 bits: ${v}`)
: p.max < 2 ** 64 ? apacheArrow.Type.Uint64
: error(`BigInt exceeds 64 bits: ${p.max}`);
}
else if (p.bools === valid) {
return apacheArrow.Type.Bool;
}
else if (p.utcdays === valid) {
return apacheArrow.Type.DateDay;
}
else if (p.dates === valid) {
return apacheArrow.Type.DateMillisecond;
}
else if (p.arrays === valid) {
const type = apacheArrow.Field.new('value', p.array_prof.type(), true);
return p.minlen === p.maxlen
? new apacheArrow.FixedSizeList(p.minlen, type)
: new apacheArrow.List(type);
}
else if (p.structs === valid) {
const sp = p.struct_prof;
return new apacheArrow.Struct(
Object.keys(sp).map(name => apacheArrow.Field.new(name, sp[name].type(), true))
);
}
else if (p.strings > 0) {
return apacheArrow.Type.Dictionary;
}
else {
error('Type inference failure');
}
}
function columnFromObjects(data, name, nrows, scan, type, nullable = true) {
type = resolveType(type);
// perform type inference if needed
if (!type) {
const p = profile(scan, name);
nullable = p.nulls > 0;
type = p.type();
}
return apacheArrow.Column.new(
name,
dataFromScan(nrows, scan, name, type, nullable)
);
}
function columnFromTable(table, name, nrows, scan, type, nullable = true) {
type = resolveType(type);
const column = table.column(name);
const reified = !(table.isFiltered() || table.isOrdered());
// use existing arrow data if types match
const vec = arrowVector$1(column);
if (vec && reified && typeCompatible(vec.type, type)) {
return apacheArrow.Column.new(name, vec.chunks || vec);
}
// if backing data is a typed array, leverage that
const data = column.data;
if (isTypedArray(data)) {
const dtype = typeFromArray(data);
if (reified && dtype && typeCompatible(dtype, type)) {
return apacheArrow.Column.new(name, dataFromArray(data, dtype));
} else {
type = type || dtype;
nullable = false;
}
}
// perform type inference if needed
if (!type) {
const p = profile(scan, column);
nullable = p.nulls > 0;
type = p.type();
}
return apacheArrow.Column.new(
name,
dataFromScan(nrows, scan, column, type, nullable)
);
}
function arrowVector$1(value) {
return value instanceof apacheArrow.Vector ? value
: value.vector instanceof apacheArrow.Vector ? value.vector
: null;
}
const types = {
Float32Array: apacheArrow.Float32,
Float64Array: apacheArrow.Float64,
Int8Array: apacheArrow.Int8,
Int16Array: apacheArrow.Int16,
Int32Array: apacheArrow.Int32,
Uint8Array: apacheArrow.Uint8,
Uint16Array: apacheArrow.Uint16,
Uint32Array: apacheArrow.Uint32,
BigInt64Array: apacheArrow.Int64,
BigUint64Array: apacheArrow.Uint64
};
function typeFromArray(data) {
const Type = types[data.constructor.name];
return Type ? new Type() : null;
}
function typeCompatible(a, b) {
return !a || !b ? true : a.compareTo(b);
}
function isFunction(value) {
return typeof value === 'function';
}
function scanArray(data, limit, offset) {
const n = Math.min(data.length, offset + limit);
return (name, visit) => {
for (let i = offset; i < n; ++i) {
visit(data[i][name], i);
}
};
}
/**
* Options for Arrow encoding.
* @typedef {object} ArrowFormatOptions
* @property {number} [limit=Infinity] The maximum number of rows to include.
* @property {number} [offset=0] The row offset indicating how many initial
* rows to skip.
* @property {string[]} [columns] Ordered list of column names to include.
* If function-valued, the function should accept a dataset as input and
* return an array of column name strings.
* @property {object} [types] The Arrow data types to use. If specified,
* theinput should be an object with column names for keys and Arrow data
* types for values. If a column type is not explicitly provided, type
* inference will be performed to guess an appropriate type.
*/
/**
* Create an Apache Arrow table for an input dataset.
* @param {Array|object} data An input dataset to convert to Arrow format.
* If array-valued, the data should consist of an array of objects where
* each entry represents a row and named properties represent columns.
* Otherwise, the input data should be an Arquero table.
* @param {ArrowFormatOptions} [options] Encoding options, including
* column data types.
* @return {Table} An Apache Arrow Table instance.
*/
function toArrow(data, options = {}) {
const types = options.types || {};
const { columnFrom, names, nrows, scan } = init(data, options);
return apacheArrow.Table.new(
names.map(name => {
const col = columnFrom(data, name, nrows, scan, types[name]);
return col.length === nrows
? col
: error('Column length mismatch');
})
);
}
function init(data, options) {
const { columns, limit = Infinity, offset = 0 } = options;
const names = isFunction(columns) ? columns(data)
: isArray(columns) ? columns
: null;
if (isArray(data)) {
return {
columnFrom: columnFromObjects,
names: names || Object.keys(data[0]),
nrows: Math.min(limit, data.length - offset),
scan: scanArray(data, limit, offset)
};
} else if (isTable(data)) {
return {
columnFrom: columnFromTable,
names: names || data.columnNames(),
nrows: Math.min(limit, data.numRows() - offset),
scan: scanTable(data, limit, offset)
};
} else {
error('Unsupported input data type');
}
}
function isTable(data) {
return data && isFunction(data.reify);
}
function dataFromTable(table, column, type, nullable) {
const nrows = table.numRows();
const scan = scanTable(table, Infinity, 0);
return dataFromScan(nrows, scan, column, type, nullable);
}
// public API
const arquero_package = {
tableMethods: { toArrow }
};
Object.defineProperty(exports, 'Type', {
enumerable: true,
get: function () {
return apacheArrow.Type;
}
});
exports.arquero_package = arquero_package;
exports.dataFromTable = dataFromTable;
exports.profiler = profiler;
exports.toArrow = toArrow;
Object.defineProperty(exports, '__esModule', { value: true });
})));