@loaders.gl/parquet 3.0.12 → 3.1.0-alpha.4

package/src/lib/parse-parquet.ts

@@ -0,0 +1,27 @@
+// import type {LoaderWithParser, Loader, LoaderOptions} from '@loaders.gl/loader-utils';
+import type {ParquetLoaderOptions} from '../parquet-loader';
+
+import {ParquetReader} from '../parquetjs/parser/parquet-reader';
+
+export async function parseParquet(arrayBuffer: ArrayBuffer, options?: ParquetLoaderOptions) {
+  const blob = new Blob([arrayBuffer]);
+  for await (const batch of parseParquetFileInBatches(blob, options)) {
+    return batch;
+  }
+  return null;
+}
+
+export async function* parseParquetFileInBatches(blob: Blob, options?: ParquetLoaderOptions) {
+  const reader = await ParquetReader.openBlob(blob);
+  const rows: any[][] = [];
+  try {
+    const cursor = reader.getCursor();
+    let record: any[] | null;
+    while ((record = await cursor.next())) {
+      rows.push(record);
+    }
+  } finally {
+    await reader.close();
+  }
+  yield rows;
+}

package/src/parquet-loader.ts

@@ -1,5 +1,4 @@
-import type {LoaderWithParser, Loader, LoaderOptions} from '@loaders.gl/loader-utils';
-import {ParquetReader} from './parquetjs/reader';
+import type {Loader, LoaderOptions} from '@loaders.gl/loader-utils';
 
 // __VERSION__ is injected by babel-plugin-version-inline
 // @ts-ignore TS2304: Cannot find name '__VERSION__'.
@@ -20,7 +19,7 @@ const DEFAULT_PARQUET_LOADER_OPTIONS: ParquetLoaderOptions = {
 };
 
 /** ParquetJS table loader */
-export const ParquetWorkerLoader: Loader = {
+export const ParquetLoader = {
   name: 'Apache Parquet',
   id: 'parquet',
   module: 'parquet',
@@ -30,27 +29,8 @@ export const ParquetWorkerLoader: Loader = {
   extensions: ['parquet'],
   mimeTypes: ['application/octet-stream'],
   binary: true,
-  // tests: ['PARQUET'],
+  tests: ['PAR1', 'PARE'],
   options: DEFAULT_PARQUET_LOADER_OPTIONS
 };
 
-/** ParquetJS table loader */
-export const ParquetLoader: LoaderWithParser = {
-  ...ParquetWorkerLoader,
-  parse
-};
-
-async function parse(arrayBuffer: ArrayBuffer, options?: ParquetLoaderOptions) {
-  const reader = await ParquetReader.openArrayBuffer(arrayBuffer);
-  const rows: any[][] = [];
-  try {
-    const cursor = reader.getCursor();
-    let record: any[] | null;
-    while ((record = await cursor.next())) {
-      rows.push(record);
-    }
-  } finally {
-    await reader.close();
-  }
-  return rows;
-}
+export const _typecheckParquetLoader: Loader = ParquetLoader;

package/src/parquetjs/codecs/dictionary.ts

@@ -0,0 +1,11 @@
+import {decodeValues as decodeRleValues} from './rle';
+
+export function decodeValues(type, cursor, count, opts) {
+  opts.bitWidth = cursor.buffer.slice(cursor.offset, cursor.offset + 1).readInt8(0);
+  cursor.offset += 1;
+  return decodeRleValues(type, cursor, count, {...opts, disableEnvelope: true});
+}
+
+export function encodeValues(type, cursor, count, opts) {
+  throw new Error('Encode dictionary functionality is not supported');
+}

package/src/parquetjs/codecs/index.ts

@@ -3,6 +3,7 @@ import type {ParquetCodec} from '../schema/declare';
 import type {ParquetCodecKit} from './declare';
 import * as PLAIN from './plain';
 import * as RLE from './rle';
+import * as DICTIONARY from './dictionary';
 
 export * from './declare';
 
@@ -14,5 +15,17 @@ export const PARQUET_CODECS: Record<ParquetCodec, ParquetCodecKit> = {
   RLE: {
     encodeValues: RLE.encodeValues,
     decodeValues: RLE.decodeValues
+  },
+  // Using the PLAIN_DICTIONARY enum value is deprecated in the Parquet 2.0 specification.
+  PLAIN_DICTIONARY: {
+    // @ts-ignore
+    encodeValues: DICTIONARY.encodeValues,
+    decodeValues: DICTIONARY.decodeValues
+  },
+  // Prefer using RLE_DICTIONARY in a data page and PLAIN in a dictionary page for Parquet 2.0+ files.
+  RLE_DICTIONARY: {
+    // @ts-ignore
+    encodeValues: DICTIONARY.encodeValues,
+    decodeValues: DICTIONARY.decodeValues
   }
 };

package/src/parquetjs/codecs/rle.ts

@@ -138,7 +138,8 @@ function decodeRunRepeated(
 
   let value = 0;
   for (let i = 0; i < Math.ceil(bitWidth / 8); i++) {
-    value <<= 8;
+    // eslint-disable-next-line
+    value << 8; //  TODO - this looks wrong
     value += cursor.buffer[cursor.offset];
     cursor.offset += 1;
   }
@@ -173,7 +174,8 @@ function encodeRunRepeated(value: number, count: number, opts: ParquetCodecOptio
 
   for (let i = 0; i < buf.length; i++) {
     buf.writeUInt8(value & 0xff, i);
-    value >>= 8;
+    // eslint-disable-next-line
+    value >> 8; //  TODO - this looks wrong
   }
 
   return Buffer.concat([Buffer.from(varint.encode(count << 1)), buf]);

package/src/parquetjs/compression.ts

@@ -1,57 +1,107 @@
 // Forked from https://github.com/kbajalc/parquets under MIT license (Copyright (c) 2017 ironSource Ltd.)
 /* eslint-disable camelcase */
+// Forked from https://github.com/ironSource/parquetjs under MIT license
+
+import {
+  Compression,
+  NoCompression,
+  GZipCompression,
+  SnappyCompression,
+  BrotliCompression,
+  LZOCompression,
+  LZ4Compression,
+  ZstdCompression
+} from '@loaders.gl/compression';
+
 import {ParquetCompression} from './schema/declare';
-import * as Util from './util';
-import zlib from 'zlib';
-import snappyjs from 'snappyjs';
+import {toArrayBuffer, toBuffer} from './utils/buffer-utils';
+
+// TODO switch to worker compression to avoid bundling...
+
+// import brotli from 'brotli'; - brotli has problems with decompress in browsers
+import brotliDecompress from 'brotli/decompress';
+import lz4js from 'lz4js';
+import lzo from 'lzo';
+import {ZstdCodec} from 'zstd-codec';
+
+// Inject large dependencies through Compression constructor options
+const modules = {
+  // brotli has problems with decompress in browsers
+  brotli: {
+    decompress: brotliDecompress,
+    compress: () => {
+      throw new Error('brotli compress');
+    }
+  },
+  lz4js,
+  lzo,
+  'zstd-codec': ZstdCodec
+};
 
-let brotli: any;
-let lzo: any;
-let lz4js: any;
+// See https://github.com/apache/parquet-format/blob/master/Compression.md
+export const PARQUET_COMPRESSION_METHODS: Record<ParquetCompression, Compression> = {
+  UNCOMPRESSED: new NoCompression(),
+  GZIP: new GZipCompression(),
+  SNAPPY: new SnappyCompression(),
+  BROTLI: new BrotliCompression({modules}),
+  // TODO: Understand difference between LZ4 and LZ4_RAW
+  LZ4: new LZ4Compression({modules}),
+  LZ4_RAW: new LZ4Compression({modules}),
+  LZO: new LZOCompression({modules}),
+  ZSTD: new ZstdCompression({modules})
+};
 
-export interface ParquetCompressionKit {
-  deflate: (value: Buffer) => Buffer;
-  inflate: (value: Buffer, size: number) => Buffer;
+/**
+ * Register compressions that have big external libraries
+ * @param options.modules External library dependencies
+ */
+export async function preloadCompressions(options?: {modules: {[key: string]: any}}) {
+  const compressions = Object.values(PARQUET_COMPRESSION_METHODS);
+  return await Promise.all(compressions.map((compression) => compression.preload()));
 }
 
-export const PARQUET_COMPRESSION_METHODS: Record<ParquetCompression, ParquetCompressionKit> = {
-  UNCOMPRESSED: {
-    deflate: deflate_identity,
-    inflate: inflate_identity
-  },
-  GZIP: {
-    deflate: deflate_gzip,
-    inflate: inflate_gzip
-  },
-  SNAPPY: {
-    deflate: deflate_snappy,
-    inflate: inflate_snappy
-  },
-  LZO: {
-    deflate: deflate_lzo,
-    inflate: inflate_lzo
-  },
-  BROTLI: {
-    deflate: deflate_brotli,
-    inflate: inflate_brotli
-  },
-  LZ4: {
-    deflate: deflate_lz4,
-    inflate: inflate_lz4
+/**
+ * Deflate a value using compression method `method`
+ */
+export async function deflate(method: ParquetCompression, value: Buffer): Promise<Buffer> {
+  const compression = PARQUET_COMPRESSION_METHODS[method];
+  if (!compression) {
+    throw new Error(`parquet: invalid compression method: ${method}`);
   }
-};
+  const inputArrayBuffer = toArrayBuffer(value);
+  const compressedArrayBuffer = await compression.compress(inputArrayBuffer);
+  return toBuffer(compressedArrayBuffer);
+}
 
 /**
- * Deflate a value using compression method `method`
+ * Inflate a value using compression method `method`
+ */
+export async function decompress(
+  method: ParquetCompression,
+  value: Buffer,
+  size: number
+): Promise<Buffer> {
+  const compression = PARQUET_COMPRESSION_METHODS[method];
+  if (!compression) {
+    throw new Error(`parquet: invalid compression method: ${method}`);
+  }
+  const inputArrayBuffer = toArrayBuffer(value);
+  const compressedArrayBuffer = await compression.decompress(inputArrayBuffer, size);
+  return toBuffer(compressedArrayBuffer);
+}
+
+/*
+ * Inflate a value using compression method `method`
  */
-export function deflate(method: ParquetCompression, value: Buffer): Buffer {
+export function inflate(method: ParquetCompression, value: Buffer, size: number): Buffer {
   if (!(method in PARQUET_COMPRESSION_METHODS)) {
     throw new Error(`invalid compression method: ${method}`);
   }
-
-  return PARQUET_COMPRESSION_METHODS[method].deflate(value);
+  // @ts-ignore
+  return PARQUET_COMPRESSION_METHODS[method].inflate(value, size);
 }
 
+/*
 function deflate_identity(value: Buffer): Buffer {
   return value;
 }
@@ -92,18 +142,6 @@ function deflate_lz4(value: Buffer): Buffer {
     throw err;
   }
 }
-
-/**
- * Inflate a value using compression method `method`
- */
-export function inflate(method: ParquetCompression, value: Buffer, size: number): Buffer {
-  if (!(method in PARQUET_COMPRESSION_METHODS)) {
-    throw new Error(`invalid compression method: ${method}`);
-  }
-
-  return PARQUET_COMPRESSION_METHODS[method].inflate(value, size);
-}
-
 function inflate_identity(value: Buffer): Buffer {
   return value;
 }
@@ -142,3 +180,4 @@ function inflate_brotli(value: Buffer): Buffer {
   }
   return Buffer.from(brotli.decompress(value));
 }
+*/

package/src/parquetjs/{writer.ts → encoder/writer.ts}

@@ -1,17 +1,17 @@
 // Forked from https://github.com/kbajalc/parquets under MIT license (Copyright (c) 2017 ironSource Ltd.)
 /* eslint-disable camelcase */
 import {Transform, Writable} from 'stream';
-import {ParquetCodecOptions, PARQUET_CODECS} from './codecs';
-import * as Compression from './compression';
+import {ParquetCodecOptions, PARQUET_CODECS} from '../codecs';
+import * as Compression from '../compression';
 import {
   ParquetBuffer,
   ParquetCodec,
   ParquetData,
   ParquetField,
   PrimitiveType
-} from './schema/declare';
-import {ParquetSchema} from './schema/schema';
-import * as Shred from './schema/shred';
+} from '../schema/declare';
+import {ParquetSchema} from '../schema/schema';
+import * as Shred from '../schema/shred';
 import {
   ColumnChunk,
   ColumnMetaData,
@@ -28,8 +28,9 @@ import {
   RowGroup,
   SchemaElement,
   Type
-} from './parquet-thrift';
-import * as Util from './util';
+} from '../parquet-thrift';
+import {osopen, oswrite, osclose} from '../utils/file-utils';
+import {getBitWidth, serializeThrift} from '../utils/read-utils';
 import Int64 from 'node-int64';
 
 /**
@@ -85,7 +86,7 @@ export class ParquetWriter<T> {
     path: string,
     opts?: ParquetWriterOptions
   ): Promise<ParquetWriter<T>> {
-    const outputStream = await Util.osopen(path, opts);
+    const outputStream = await osopen(path, opts);
     return ParquetWriter.openStream(schema, outputStream, opts);
   }
 
@@ -131,13 +132,16 @@ export class ParquetWriter<T> {
     this.closed = false;
     this.userMetadata = {};
 
+    // eslint-disable-next-line @typescript-eslint/no-floating-promises
+    this.writeHeader();
+  }
+
+  async writeHeader(): Promise<void> {
     // TODO - better not mess with promises in the constructor
     try {
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      envelopeWriter.writeHeader();
+      await this.envelopeWriter.writeHeader();
     } catch (err) {
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      envelopeWriter.close();
+      await this.envelopeWriter.close();
       throw err;
     }
   }
@@ -226,8 +230,8 @@ export class ParquetEnvelopeWriter {
     outputStream: Writable,
     opts: ParquetWriterOptions
   ): Promise<ParquetEnvelopeWriter> {
-    const writeFn = Util.oswrite.bind(undefined, outputStream);
-    const closeFn = Util.osclose.bind(undefined, outputStream);
+    const writeFn = oswrite.bind(undefined, outputStream);
+    const closeFn = osclose.bind(undefined, outputStream);
     return new ParquetEnvelopeWriter(schema, writeFn, closeFn, 0, opts);
   }
 
@@ -274,7 +278,7 @@ export class ParquetEnvelopeWriter {
    * shredRecord method
    */
   async writeRowGroup(records: ParquetBuffer): Promise<void> {
-    const rgroup = encodeRowGroup(this.schema, records, {
+    const rgroup = await encodeRowGroup(this.schema, records, {
       baseOffset: this.offset,
       pageSize: this.pageSize,
       useDataPageV2: this.useDataPageV2
@@ -282,7 +286,7 @@ export class ParquetEnvelopeWriter {
 
     this.rowCount += records.rowCount;
     this.rowGroups.push(rgroup.metadata);
-    return this.writeSection(rgroup.body);
+    return await this.writeSection(rgroup.body);
   }
 
   /**
@@ -363,19 +367,19 @@ function encodeValues(
 /**
  * Encode a parquet data page
  */
-function encodeDataPage(
+async function encodeDataPage(
   column: ParquetField,
   data: ParquetData
-): {
+): Promise<{
   header: PageHeader;
   headerSize: number;
   page: Buffer;
-} {
+}> {
   /* encode repetition and definition levels */
   let rLevelsBuf = Buffer.alloc(0);
   if (column.rLevelMax > 0) {
     rLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.rlevels, {
-      bitWidth: Util.getBitWidth(column.rLevelMax)
+      bitWidth: getBitWidth(column.rLevelMax)
       // disableEnvelope: false
     });
   }
@@ -383,7 +387,7 @@ function encodeDataPage(
   let dLevelsBuf = Buffer.alloc(0);
   if (column.dLevelMax > 0) {
     dLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.dlevels, {
-      bitWidth: Util.getBitWidth(column.dLevelMax)
+      bitWidth: getBitWidth(column.dLevelMax)
       // disableEnvelope: false
     });
   }
@@ -397,7 +401,7 @@ function encodeDataPage(
   const dataBuf = Buffer.concat([rLevelsBuf, dLevelsBuf, valuesBuf]);
 
   // compression = column.compression === 'UNCOMPRESSED' ? (compression || 'UNCOMPRESSED') : column.compression;
-  const compressedBuf = Compression.deflate(column.compression!, dataBuf);
+  const compressedBuf = await Compression.deflate(column.compression!, dataBuf);
 
   /* build page header */
   const header = new PageHeader({
@@ -413,7 +417,7 @@ function encodeDataPage(
   });
 
   /* concat page header, repetition and definition levels and values */
-  const headerBuf = Util.serializeThrift(header);
+  const headerBuf = serializeThrift(header);
   const page = Buffer.concat([headerBuf, compressedBuf]);
 
   return {header, headerSize: headerBuf.length, page};
@@ -422,15 +426,15 @@ function encodeDataPage(
 /**
  * Encode a parquet data page (v2)
  */
-function encodeDataPageV2(
+async function encodeDataPageV2(
   column: ParquetField,
   data: ParquetData,
   rowCount: number
-): {
+): Promise<{
   header: PageHeader;
   headerSize: number;
   page: Buffer;
-} {
+}> {
   /* encode values */
   const valuesBuf = encodeValues(column.primitiveType!, column.encoding!, data.values, {
     typeLength: column.typeLength,
@@ -438,13 +442,13 @@ function encodeDataPageV2(
   });
 
   // compression = column.compression === 'UNCOMPRESSED' ? (compression || 'UNCOMPRESSED') : column.compression;
-  const compressedBuf = Compression.deflate(column.compression!, valuesBuf);
+  const compressedBuf = await Compression.deflate(column.compression!, valuesBuf);
 
   /* encode repetition and definition levels */
   let rLevelsBuf = Buffer.alloc(0);
   if (column.rLevelMax > 0) {
     rLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.rlevels, {
-      bitWidth: Util.getBitWidth(column.rLevelMax),
+      bitWidth: getBitWidth(column.rLevelMax),
       disableEnvelope: true
     });
   }
@@ -452,7 +456,7 @@ function encodeDataPageV2(
   let dLevelsBuf = Buffer.alloc(0);
   if (column.dLevelMax > 0) {
     dLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.dlevels, {
-      bitWidth: Util.getBitWidth(column.dLevelMax),
+      bitWidth: getBitWidth(column.dLevelMax),
       disableEnvelope: true
     });
   }
@@ -474,7 +478,7 @@ function encodeDataPageV2(
   });
 
   /* concat page header, repetition and definition levels and values */
-  const headerBuf = Util.serializeThrift(header);
+  const headerBuf = serializeThrift(header);
   const page = Buffer.concat([headerBuf, rLevelsBuf, dLevelsBuf, compressedBuf]);
   return {header, headerSize: headerBuf.length, page};
 }
@@ -482,16 +486,16 @@ function encodeDataPageV2(
 /**
  * Encode an array of values into a parquet column chunk
  */
-function encodeColumnChunk(
+async function encodeColumnChunk(
   column: ParquetField,
   buffer: ParquetBuffer,
   offset: number,
   opts: ParquetWriterOptions
-): {
+): Promise<{
   body: Buffer;
   metadata: ColumnMetaData;
   metadataOffset: number;
-} {
+}> {
   const data = buffer.columnData[column.path.join()];
   const baseOffset = (opts.baseOffset || 0) + offset;
   /* encode data page(s) */
@@ -502,12 +506,9 @@ function encodeColumnChunk(
   // tslint:disable-next-line:variable-name
   let total_compressed_size = 0;
   {
-    let result: any;
-    if (opts.useDataPageV2) {
-      result = encodeDataPageV2(column, data, buffer.rowCount);
-    } else {
-      result = encodeDataPage(column, data);
-    }
+    const result = opts.useDataPageV2
+      ? await encodeDataPageV2(column, data, buffer.rowCount)
+      : await encodeDataPage(column, data);
     // pages.push(result.page);
     pageBuf = result.page;
     total_uncompressed_size += result.header.uncompressed_page_size + result.headerSize;
@@ -535,21 +536,21 @@ function encodeColumnChunk(
 
   /* concat metadata header and data pages */
   const metadataOffset = baseOffset + pageBuf.length;
-  const body = Buffer.concat([pageBuf, Util.serializeThrift(metadata)]);
+  const body = Buffer.concat([pageBuf, serializeThrift(metadata)]);
   return {body, metadata, metadataOffset};
 }
 
 /**
  * Encode a list of column values into a parquet row group
  */
-function encodeRowGroup(
+async function encodeRowGroup(
   schema: ParquetSchema,
   data: ParquetBuffer,
   opts: ParquetWriterOptions
-): {
+): Promise<{
   body: Buffer;
   metadata: RowGroup;
-} {
+}> {
   const metadata = new RowGroup({
     num_rows: data.rowCount,
     columns: [],
@@ -562,7 +563,7 @@ function encodeRowGroup(
       continue; // eslint-disable-line no-continue
     }
 
-    const cchunkData = encodeColumnChunk(field, data, body.length, opts);
+    const cchunkData = await encodeColumnChunk(field, data, body.length, opts);
 
     const cchunk = new ColumnChunk({
       file_offset: cchunkData.metadataOffset,
@@ -634,7 +635,7 @@ function encodeFooter(
     metadata.schema.push(schemaElem);
   }
 
-  const metadataEncoded = Util.serializeThrift(metadata);
+  const metadataEncoded = serializeThrift(metadata);
   const footerEncoded = Buffer.alloc(metadataEncoded.length + 8);
   metadataEncoded.copy(footerEncoded);
   footerEncoded.writeUInt32LE(metadataEncoded.length, metadataEncoded.length);