@loaders.gl/parquet 4.4.5 → 5.0.0-alpha.2

package/src/parquetjs/codecs/plain.ts

@@ -6,14 +6,29 @@
 
 /* eslint-disable camelcase */
 import type {PrimitiveType} from '../schema/declare';
-import type {CursorBuffer, ParquetCodecOptions} from './declare';
-import INT53 from 'int53';
+import {
+  getParquetValueOutput,
+  type CursorBuffer,
+  type ParquetCodecOptions,
+  type ParquetValueBuffer
+} from './declare';
+import {
+  concatUint8Arrays,
+  copyUint8Array,
+  encodeUtf8,
+  toUint8Array,
+  writeDoubleLE,
+  writeFloatLE,
+  writeInt32LE,
+  writeInt64LE,
+  writeUInt32LE
+} from '../utils/binary-utils';
 
 export function encodeValues(
   type: PrimitiveType,
   values: any[],
   opts: ParquetCodecOptions
-): Buffer {
+): Uint8Array {
   switch (type) {
     case 'BOOLEAN':
       return encodeValues_BOOLEAN(values);
@@ -41,22 +56,22 @@ export function decodeValues(
   cursor: CursorBuffer,
   count: number,
   opts: ParquetCodecOptions
-): any[] {
+): ParquetValueBuffer {
   switch (type) {
     case 'BOOLEAN':
-      return decodeValues_BOOLEAN(cursor, count);
+      return decodeValues_BOOLEAN(cursor, count, opts);
     case 'INT32':
-      return decodeValues_INT32(cursor, count);
+      return decodeValues_INT32(cursor, count, opts);
     case 'INT64':
-      return decodeValues_INT64(cursor, count);
+      return decodeValues_INT64(cursor, count, opts);
     case 'INT96':
-      return decodeValues_INT96(cursor, count);
+      return decodeValues_INT96(cursor, count, opts);
     case 'FLOAT':
-      return decodeValues_FLOAT(cursor, count);
+      return decodeValues_FLOAT(cursor, count, opts);
     case 'DOUBLE':
-      return decodeValues_DOUBLE(cursor, count);
+      return decodeValues_DOUBLE(cursor, count, opts);
     case 'BYTE_ARRAY':
-      return decodeValues_BYTE_ARRAY(cursor, count);
+      return decodeValues_BYTE_ARRAY(cursor, count, opts);
     case 'FIXED_LEN_BYTE_ARRAY':
       return decodeValues_FIXED_LEN_BYTE_ARRAY(cursor, count, opts);
     default:
@@ -64,9 +79,8 @@ export function decodeValues(
   }
 }
 
-function encodeValues_BOOLEAN(values: boolean[]): Buffer {
-  const buf = Buffer.alloc(Math.ceil(values.length / 8));
-  buf.fill(0);
+function encodeValues_BOOLEAN(values: boolean[]): Uint8Array {
+  const buf = new Uint8Array(Math.ceil(values.length / 8));
   for (let i = 0; i < values.length; i++) {
     if (values[i]) {
       buf[Math.floor(i / 8)] |= 1 << (i % 8);
@@ -75,168 +89,216 @@ function encodeValues_BOOLEAN(values: boolean[]): Buffer {
   return buf;
 }
 
-function decodeValues_BOOLEAN(cursor: CursorBuffer, count: number): boolean[] {
-  const values: boolean[] = [];
+function decodeValues_BOOLEAN(
+  cursor: CursorBuffer,
+  count: number,
+  options: ParquetCodecOptions
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(options, count);
   for (let i = 0; i < count; i++) {
     const b = cursor.buffer[cursor.offset + Math.floor(i / 8)];
-    values.push((b & (1 << (i % 8))) > 0);
+    const value = (b & (1 << (i % 8))) > 0;
+    output[outputOffset + i] = output instanceof Uint8Array ? Number(value) : value;
   }
   cursor.offset += Math.ceil(count / 8);
-  return values;
+  return output;
 }
 
-function encodeValues_INT32(values: number[]): Buffer {
-  const buf = Buffer.alloc(4 * values.length);
+function encodeValues_INT32(values: number[]): Uint8Array {
+  const buf = new Uint8Array(4 * values.length);
   for (let i = 0; i < values.length; i++) {
-    buf.writeInt32LE(values[i], i * 4);
+    writeInt32LE(buf, values[i], i * 4);
   }
   return buf;
 }
 
-function decodeValues_INT32(cursor: CursorBuffer, count: number): number[] {
-  const values: number[] = [];
+function decodeValues_INT32(
+  cursor: CursorBuffer,
+  count: number,
+  options: ParquetCodecOptions
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(options, count);
+  const dataView = getCursorDataView(cursor);
   for (let i = 0; i < count; i++) {
-    values.push(cursor.buffer.readInt32LE(cursor.offset));
+    output[outputOffset + i] = dataView.getInt32(cursor.offset, true);
     cursor.offset += 4;
   }
-  return values;
+  return output;
 }
 
-function encodeValues_INT64(values: number[]): Buffer {
-  const buf = Buffer.alloc(8 * values.length);
+function encodeValues_INT64(values: Array<number | bigint>): Uint8Array {
+  const buf = new Uint8Array(8 * values.length);
   for (let i = 0; i < values.length; i++) {
-    INT53.writeInt64LE(values[i], buf, i * 8);
+    writeInt64LE(buf, values[i], i * 8);
   }
   return buf;
 }
 
-function decodeValues_INT64(cursor: CursorBuffer, count: number): number[] {
-  const values: number[] = [];
+function decodeValues_INT64(
+  cursor: CursorBuffer,
+  count: number,
+  options: ParquetCodecOptions
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(options, count);
+  const dataView = getCursorDataView(cursor);
   for (let i = 0; i < count; i++) {
-    values.push(INT53.readInt64LE(cursor.buffer, cursor.offset));
+    const value = dataView.getBigInt64(cursor.offset, true);
+    output[outputOffset + i] = options.int64AsBigInt ? value : Number(value);
     cursor.offset += 8;
   }
-  return values;
+  return output;
 }
 
-function encodeValues_INT96(values: number[]): Buffer {
-  const buf = Buffer.alloc(12 * values.length);
+function encodeValues_INT96(values: number[]): Uint8Array {
+  const buf = new Uint8Array(12 * values.length);
   for (let i = 0; i < values.length; i++) {
-    if (values[i] >= 0) {
-      INT53.writeInt64LE(values[i], buf, i * 12);
-      buf.writeUInt32LE(0, i * 12 + 8); // truncate to 64 actual precision
-    } else {
-      INT53.writeInt64LE(~-values[i] + 1, buf, i * 12);
-      buf.writeUInt32LE(0xffffffff, i * 12 + 8); // truncate to 64 actual precision
-    }
+    writeInt64LE(buf, values[i], i * 12);
+    writeUInt32LE(buf, values[i] >= 0 ? 0 : 0xffffffff, i * 12 + 8);
   }
   return buf;
 }
 
-function decodeValues_INT96(cursor: CursorBuffer, count: number): number[] {
-  const values: number[] = [];
+function decodeValues_INT96(
+  cursor: CursorBuffer,
+  count: number,
+  options: ParquetCodecOptions
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(options, count);
+  const dataView = getCursorDataView(cursor);
   for (let i = 0; i < count; i++) {
-    const low = INT53.readInt64LE(cursor.buffer, cursor.offset);
-    const high = cursor.buffer.readUInt32LE(cursor.offset + 8);
+    const low = Number(dataView.getBigInt64(cursor.offset, true));
+    const high = dataView.getUint32(cursor.offset + 8, true);
     if (high === 0xffffffff) {
-      values.push(~-low + 1); // truncate to 64 actual precision
+      output[outputOffset + i] = ~-low + 1; // truncate to 64 actual precision
     } else {
-      values.push(low); // truncate to 64 actual precision
+      output[outputOffset + i] = low; // truncate to 64 actual precision
     }
     cursor.offset += 12;
   }
-  return values;
+  return output;
 }
 
-function encodeValues_FLOAT(values: number[]): Buffer {
-  const buf = Buffer.alloc(4 * values.length);
+function encodeValues_FLOAT(values: number[]): Uint8Array {
+  const buf = new Uint8Array(4 * values.length);
   for (let i = 0; i < values.length; i++) {
-    buf.writeFloatLE(values[i], i * 4);
+    writeFloatLE(buf, values[i], i * 4);
   }
   return buf;
 }
 
-function decodeValues_FLOAT(cursor: CursorBuffer, count: number): number[] {
-  const values: number[] = [];
+function decodeValues_FLOAT(
+  cursor: CursorBuffer,
+  count: number,
+  options: ParquetCodecOptions
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(options, count);
+  const dataView = getCursorDataView(cursor);
   for (let i = 0; i < count; i++) {
-    values.push(cursor.buffer.readFloatLE(cursor.offset));
+    output[outputOffset + i] = dataView.getFloat32(cursor.offset, true);
     cursor.offset += 4;
   }
-  return values;
+  return output;
 }
 
-function encodeValues_DOUBLE(values: number[]): Buffer {
-  const buf = Buffer.alloc(8 * values.length);
+function encodeValues_DOUBLE(values: number[]): Uint8Array {
+  const buf = new Uint8Array(8 * values.length);
   for (let i = 0; i < values.length; i++) {
-    buf.writeDoubleLE(values[i], i * 8);
+    writeDoubleLE(buf, values[i], i * 8);
   }
   return buf;
 }
 
-function decodeValues_DOUBLE(cursor: CursorBuffer, count: number): number[] {
-  const values: number[] = [];
+function decodeValues_DOUBLE(
+  cursor: CursorBuffer,
+  count: number,
+  options: ParquetCodecOptions
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(options, count);
+  const dataView = getCursorDataView(cursor);
   for (let i = 0; i < count; i++) {
-    values.push(cursor.buffer.readDoubleLE(cursor.offset));
+    output[outputOffset + i] = dataView.getFloat64(cursor.offset, true);
     cursor.offset += 8;
   }
-  return values;
+  return output;
 }
 
-function encodeValues_BYTE_ARRAY(values: Buffer[]): Buffer {
+function encodeValues_BYTE_ARRAY(values: any[]): Uint8Array {
+  const byteValues = values.map(toPrimitiveBytes);
   // tslint:disable-next-line:variable-name
   let buf_len = 0;
-  for (let i = 0; i < values.length; i++) {
-    values[i] = Buffer.from(values[i]);
-    buf_len += 4 + values[i].length;
+  for (let i = 0; i < byteValues.length; i++) {
+    buf_len += 4 + byteValues[i].length;
   }
-  const buf = Buffer.alloc(buf_len);
+  const buf = new Uint8Array(buf_len);
   // tslint:disable-next-line:variable-name
   let buf_pos = 0;
-  for (let i = 0; i < values.length; i++) {
-    buf.writeUInt32LE(values[i].length, buf_pos);
-    values[i].copy(buf as Uint8Array, buf_pos + 4);
-    buf_pos += 4 + values[i].length;
+  for (let i = 0; i < byteValues.length; i++) {
+    writeUInt32LE(buf, byteValues[i].length, buf_pos);
+    buf.set(byteValues[i], buf_pos + 4);
+    buf_pos += 4 + byteValues[i].length;
   }
   return buf;
 }
 
-function decodeValues_BYTE_ARRAY(cursor: CursorBuffer, count: number): Buffer[] {
-  const values: Buffer[] = [];
+function decodeValues_BYTE_ARRAY(
+  cursor: CursorBuffer,
+  count: number,
+  options: ParquetCodecOptions
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(options, count);
+  const dataView = getCursorDataView(cursor);
   for (let i = 0; i < count; i++) {
-    const len = cursor.buffer.readUInt32LE(cursor.offset);
+    const len = dataView.getUint32(cursor.offset, true);
     cursor.offset += 4;
-    // values.push(cursor.buffer.buffer.slice(cursor.offset, cursor.offset + len));
-    values.push(cursor.buffer.slice(cursor.offset, cursor.offset + len));
-    cursor.offset += len;
+    output[outputOffset + i] = readByteArray(cursor, len, options.retainByteArrayViews);
   }
-  return values;
+  return output;
 }
 
-function encodeValues_FIXED_LEN_BYTE_ARRAY(values: Buffer[], opts: ParquetCodecOptions): Buffer {
+function encodeValues_FIXED_LEN_BYTE_ARRAY(values: any[], opts: ParquetCodecOptions): Uint8Array {
   if (!opts.typeLength) {
     throw new Error('missing option: typeLength (required for FIXED_LEN_BYTE_ARRAY)');
   }
+  const byteValues = values.map(toPrimitiveBytes);
   for (let i = 0; i < values.length; i++) {
-    values[i] = Buffer.from(values[i]);
-    if (values[i].length !== opts.typeLength) {
+    if (byteValues[i].length !== opts.typeLength) {
       throw new Error(`invalid value for FIXED_LEN_BYTE_ARRAY: ${values[i]}`);
     }
   }
-  return Buffer.concat(values as Uint8Array[]);
+  return concatUint8Arrays(byteValues);
 }
 
 function decodeValues_FIXED_LEN_BYTE_ARRAY(
   cursor: CursorBuffer,
   count: number,
   opts: ParquetCodecOptions
-): Buffer[] {
-  const values: Buffer[] = [];
+): ParquetValueBuffer {
+  const {output, outputOffset} = getParquetValueOutput(opts, count);
   if (!opts.typeLength) {
     throw new Error('missing option: typeLength (required for FIXED_LEN_BYTE_ARRAY)');
   }
   for (let i = 0; i < count; i++) {
-    values.push(cursor.buffer.slice(cursor.offset, cursor.offset + opts.typeLength));
-    cursor.offset += opts.typeLength;
+    output[outputOffset + i] = readByteArray(cursor, opts.typeLength, opts.retainByteArrayViews);
   }
-  return values;
+  return output;
+}
+
+/** Reads one byte-array value, optionally retaining a view into the decoded page buffer. */
+function readByteArray(
+  cursor: CursorBuffer,
+  byteLength: number,
+  retainByteArrayView: boolean | undefined
+): Uint8Array {
+  const value = cursor.buffer.subarray(cursor.offset, cursor.offset + byteLength);
+  cursor.offset += byteLength;
+  return retainByteArrayView ? value : copyUint8Array(value);
+}
+
+/** Creates one reusable DataView for all primitive reads from a codec cursor. */
+function getCursorDataView(cursor: CursorBuffer): DataView {
+  return new DataView(cursor.buffer.buffer, cursor.buffer.byteOffset, cursor.buffer.byteLength);
+}
+
+function toPrimitiveBytes(value: any): Uint8Array {
+  return typeof value === 'string' ? encodeUtf8(value) : toUint8Array(value);
 }