arquero-arrow/dist/arquero-arrow.mjs

Arrow serialization support for Arquero.

import { Data, Vector, Builder, Type, DataType, Utf8, Uint64, Uint32, Uint16, Uint8, TimeSecond, TimeNanosecond, TimeMillisecond, TimeMicrosecond, Null, IntervalYearMonth, IntervalDayTime, Int64, Int32, Int16, Int8, Float64, Float32, Float16, Dictionary, DateMillisecond, DateDay, Bool, Binary, Field, FixedSizeList, List, Struct, Column, Table } from 'apache-arrow';
export { Type } from 'apache-arrow';

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 Data
    ? d
    : Data.new(d.type, 0, d.length, d.nulls, d.buffers, null, d.dict);
}

function arrowVector(data) {
  return 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 = 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 Type.Int:
      method = type.bitWidth < 64 ? arrayBuilder : null;
      break;
    case Type.Float:
      method = type.precision > 0 ? arrayBuilder : null;
      break;
    case Type.Dictionary:
      // check sub-types against builder assumptions
      // if check fails, fallback to default builder
      method = (
        type.dictionary.typeId === Type.Utf8 &&
        type.indices.typeId === Type.Int &&
        type.indices.bitWidth < 64
      ) ? dictionaryBuilder : null;
      break;
    case Type.Bool:
      method = boolBuilder;
      break;
    case 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 DataType || type == null) {
    return type;
  }

  switch (type) {
    case Type.Binary:
      return new Binary();
    case Type.Bool:
      return new Bool();
    case Type.DateDay:
      return new DateDay();
    case Type.DateMillisecond:
    case Type.Date:
      return new DateMillisecond();
    case Type.Dictionary:
      return new Dictionary(new Utf8(), new Int32());
    case Type.Float16:
      return new Float16();
    case Type.Float32:
      return new Float32();
    case Type.Float64:
    case Type.Float:
      return new Float64();
    case Type.Int8:
      return new Int8();
    case Type.Int16:
      return new Int16();
    case Type.Int32:
    case Type.Int:
      return new Int32();
    case Type.Int64:
      return new Int64();
    case Type.IntervalDayTime:
      return new IntervalDayTime();
    case Type.Interval:
    case Type.IntervalYearMonth:
      return new IntervalYearMonth();
    case Type.Null:
      return new Null();
    case Type.TimeMicrosecond:
      return new TimeMicrosecond();
    case Type.TimeMillisecond:
    case Type.Time:
      return new TimeMillisecond();
    case Type.TimeNanosecond:
      return new TimeNanosecond();
    case Type.TimeSecond:
      return new TimeSecond();
    case Type.Uint8:
      return new Uint8();
    case Type.Uint16:
      return new Uint16();
    case Type.Uint32:
      return new Uint32();
    case Type.Uint64:
      return new Uint64();
    case Type.Utf8:
      return new 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 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 ? Type.Float64
        : v < (1 << 7) ? Type.Int8 : v < (1 << 15) ? Type.Int16 : Type.Int32
      : v >= 2 ** 32 ? Type.Float64
        : v < (1 << 8) ? Type.Uint8 : v < (1 << 16) ? Type.Uint16 : Type.Uint32;
  }
  else if (p.nums === valid) {
    return 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 ? Type.Int64
        : error(`BigInt exceeds 64 bits: ${v}`)
      : p.max < 2 ** 64 ? Type.Uint64
        : error(`BigInt exceeds 64 bits: ${p.max}`);
  }
  else if (p.bools === valid) {
    return Type.Bool;
  }
  else if (p.utcdays === valid) {
    return Type.DateDay;
  }
  else if (p.dates === valid) {
    return Type.DateMillisecond;
  }
  else if (p.arrays === valid) {
    const type = Field.new('value', p.array_prof.type(), true);
    return p.minlen === p.maxlen
      ? new FixedSizeList(p.minlen, type)
      : new List(type);
  }
  else if (p.structs === valid) {
    const sp = p.struct_prof;
    return new Struct(
      Object.keys(sp).map(name => Field.new(name, sp[name].type(), true))
    );
  }
  else if (p.strings > 0) {
    return 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 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 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 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 Column.new(
    name,
    dataFromScan(nrows, scan, column, type, nullable)
  );
}

function arrowVector$1(value) {
  return value instanceof Vector ? value
    : value.vector instanceof Vector ? value.vector
    : null;
}

const types = {
  Float32Array:    Float32,
  Float64Array:    Float64,
  Int8Array:       Int8,
  Int16Array:      Int16,
  Int32Array:      Int32,
  Uint8Array:      Uint8,
  Uint16Array:     Uint16,
  Uint32Array:     Uint32,
  BigInt64Array:   Int64,
  BigUint64Array:  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 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 }
};

export { arquero_package, dataFromTable, profiler, toArrow };