@loaders.gl/parquet 3.3.0-alpha.8 → 3.3.0

package/src/lib/parse-parquet.ts

@@ -1,6 +1,7 @@
 // import type {LoaderWithParser, Loader, LoaderOptions} from '@loaders.gl/loader-utils';
+// import {ColumnarTableBatch} from '@loaders.gl/schema';
+import {makeReadableFile} 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) {
@@ -12,16 +13,28 @@ export async function parseParquet(arrayBuffer: ArrayBuffer, options?: ParquetLo
 }
 
 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();
+  const file = makeReadableFile(blob);
+  const reader = new ParquetReader(file);
+  const rowBatches = reader.rowBatchIterator(options?.parquet);
+  for await (const rows of rowBatches) {
+    yield rows;
   }
-  yield rows;
 }
+
+// export async function* parseParquetFileInColumnarBatches(blob: Blob, options?: {columnList?: string[][]}): AsyncIterable<ColumnarTableBatch> {
+//   const rowGroupReader = new ParquetRowGroupReader({data: blob, columnList: options?.columnList});
+//   try {
+//     for await (const rowGroup of rowGroupReader) {
+//       yield convertRowGroupToTableBatch(rowGroup);
+//     }
+//   } finally {
+//     await rowGroupReader.close();
+//   }
+// }
+
+// function convertRowGroupToTableBatch(rowGroup): ColumnarTableBatch {
+//   // @ts-expect-error
+//   return {
+//     data: rowGroup
+//   };
+// }

package/src/parquet-loader.ts

@@ -8,13 +8,15 @@ export type ParquetLoaderOptions = LoaderOptions & {
   parquet?: {
     type?: 'object-row-table';
     url?: string;
+    columnList?: string[] | string[][];
   };
 };
 
 const DEFAULT_PARQUET_LOADER_OPTIONS: ParquetLoaderOptions = {
   parquet: {
     type: 'object-row-table',
-    url: undefined
+    url: undefined,
+    columnList: []
   }
 };
 

package/src/parquetjs/compression.ts

@@ -14,7 +14,20 @@ import {
 } from '@loaders.gl/compression';
 
 import {ParquetCompression} from './schema/declare';
-import {toArrayBuffer, toBuffer} from './utils/buffer-utils';
+
+/** We can't use loaders-util buffer handling since we are dependent on buffers even in the browser */
+function toBuffer(arrayBuffer: ArrayBuffer): Buffer {
+  return Buffer.from(arrayBuffer);
+}
+
+function toArrayBuffer(buffer: Buffer): ArrayBuffer {
+  // TODO - per docs we should just be able to call buffer.buffer, but there are issues
+  if (Buffer.isBuffer(buffer)) {
+    const typedArray = new Uint8Array(buffer.buffer, buffer.byteOffset, buffer.length);
+    return typedArray.slice().buffer;
+  }
+  return buffer;
+}
 
 // TODO switch to worker compression to avoid bundling...
 

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

@@ -1,6 +1,6 @@
 // Forked from https://github.com/kbajalc/parquets under MIT license (Copyright (c) 2017 ironSource Ltd.)
 /* eslint-disable camelcase */
-import {Transform, Writable} from 'stream';
+import {stream} from '@loaders.gl/loader-utils';
 import {ParquetCodecOptions, PARQUET_CODECS} from '../codecs';
 import * as Compression from '../compression';
 import {
@@ -55,7 +55,7 @@ const PARQUET_DEFAULT_ROW_GROUP_SIZE = 4096;
 const PARQUET_RDLVL_TYPE = 'INT32';
 const PARQUET_RDLVL_ENCODING = 'RLE';
 
-export interface ParquetWriterOptions {
+export interface ParquetEncoderOptions {
   baseOffset?: number;
   rowGroupSize?: number;
   pageSize?: number;
@@ -71,12 +71,12 @@ export interface ParquetWriterOptions {
 }
 
 /**
- * Write a parquet file to an output stream. The ParquetWriter will perform
+ * Write a parquet file to an output stream. The ParquetEncoder will perform
  * buffering/batching for performance, so close() must be called after all rows
  * are written.
  */
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
-export class ParquetWriter<T> {
+export class ParquetEncoder<T> {
   /**
    * Convenience method to create a new buffered parquet writer that writes to
    * the specified file
@@ -84,10 +84,10 @@ export class ParquetWriter<T> {
   static async openFile<T>(
     schema: ParquetSchema,
     path: string,
-    opts?: ParquetWriterOptions
-  ): Promise<ParquetWriter<T>> {
+    opts?: ParquetEncoderOptions
+  ): Promise<ParquetEncoder<T>> {
     const outputStream = await osopen(path, opts);
-    return ParquetWriter.openStream(schema, outputStream, opts);
+    return ParquetEncoder.openStream(schema, outputStream, opts);
   }
 
   /**
@@ -96,17 +96,11 @@ export class ParquetWriter<T> {
    */
   static async openStream<T>(
     schema: ParquetSchema,
-    outputStream: Writable,
-    opts?: ParquetWriterOptions
-  ): Promise<ParquetWriter<T>> {
-    if (!opts) {
-      // tslint:disable-next-line:no-parameter-reassignment
-      opts = {};
-    }
-
+    outputStream: stream.Writable,
+    opts: ParquetEncoderOptions = {}
+  ): Promise<ParquetEncoder<T>> {
     const envelopeWriter = await ParquetEnvelopeWriter.openStream(schema, outputStream, opts);
-
-    return new ParquetWriter(schema, envelopeWriter, opts);
+    return new ParquetEncoder(schema, envelopeWriter, opts);
   }
 
   public schema: ParquetSchema;
@@ -122,7 +116,7 @@ export class ParquetWriter<T> {
   constructor(
     schema: ParquetSchema,
     envelopeWriter: ParquetEnvelopeWriter,
-    opts: ParquetWriterOptions
+    opts: ParquetEncoderOptions
   ) {
     this.schema = schema;
     this.envelopeWriter = envelopeWriter;
@@ -227,8 +221,8 @@ export class ParquetEnvelopeWriter {
    */
   static async openStream(
     schema: ParquetSchema,
-    outputStream: Writable,
-    opts: ParquetWriterOptions
+    outputStream: stream.Writable,
+    opts: ParquetEncoderOptions
   ): Promise<ParquetEnvelopeWriter> {
     const writeFn = oswrite.bind(undefined, outputStream);
     const closeFn = osclose.bind(undefined, outputStream);
@@ -249,7 +243,7 @@ export class ParquetEnvelopeWriter {
     writeFn: (buf: Buffer) => Promise<void>,
     closeFn: () => Promise<void>,
     fileOffset: number,
-    opts: ParquetWriterOptions
+    opts: ParquetEncoderOptions
   ) {
     this.schema = schema;
     this.write = writeFn;
@@ -314,11 +308,10 @@ export class ParquetEnvelopeWriter {
 
 /**
  * Create a parquet transform stream
- */
-export class ParquetTransformer<T> extends Transform {
-  public writer: ParquetWriter<T>;
+export class ParquetTransformer<T> extends stream.Transform {
+  public writer: ParquetEncoder<T>;
 
-  constructor(schema: ParquetSchema, opts: ParquetWriterOptions = {}) {
+  constructor(schema: ParquetSchema, opts: ParquetEncoderOptions = {}) {
     super({objectMode: true});
 
     const writeProxy = (function (t: ParquetTransformer<any>) {
@@ -327,7 +320,7 @@ export class ParquetTransformer<T> extends Transform {
       };
     })(this);
 
-    this.writer = new ParquetWriter(
+    this.writer = new ParquetEncoder(
       schema,
       new ParquetEnvelopeWriter(schema, writeProxy, async () => {}, 0, opts),
       opts
@@ -348,6 +341,7 @@ export class ParquetTransformer<T> extends Transform {
     await this.writer.close(callback);
   }
 }
+ */
 
 /**
  * Encode a consecutive array of data using one of the parquet encodings
@@ -490,7 +484,7 @@ async function encodeColumnChunk(
   column: ParquetField,
   buffer: ParquetBuffer,
   offset: number,
-  opts: ParquetWriterOptions
+  opts: ParquetEncoderOptions
 ): Promise<{
   body: Buffer;
   metadata: ColumnMetaData;
@@ -546,7 +540,7 @@ async function encodeColumnChunk(
 async function encodeRowGroup(
   schema: ParquetSchema,
   data: ParquetBuffer,
-  opts: ParquetWriterOptions
+  opts: ParquetEncoderOptions
 ): Promise<{
   body: Buffer;
   metadata: RowGroup;

package/src/parquetjs/parser/parquet-reader.ts

@@ -1,163 +1,280 @@
 // Forked from https://github.com/kbajalc/parquets under MIT license (Copyright (c) 2017 ironSource Ltd.)
-import {ParquetEnvelopeReader} from './parquet-envelope-reader';
-import {FileMetaData} from '../parquet-thrift';
+import type {ReadableFile} from '@loaders.gl/loader-utils';
+
 import {ParquetSchema} from '../schema/schema';
-import {ParquetCursor} from './parquet-cursor';
-import {PARQUET_VERSION} from '../../constants';
 import {decodeSchema} from './decoders';
+import {materializeRecords} from '../schema/shred';
+
+import {PARQUET_MAGIC, PARQUET_MAGIC_ENCRYPTED} from '../../constants';
+import {ColumnChunk, CompressionCodec, FileMetaData, RowGroup, Type} from '../parquet-thrift';
+import {
+  ParquetBuffer,
+  ParquetCompression,
+  ParquetData,
+  PrimitiveType,
+  ParquetOptions
+} from '../schema/declare';
+import {decodeFileMetadata, getThriftEnum, fieldIndexOf} from '../utils/read-utils';
+import {decodeDataPages, decodePage} from './decoders';
+
+export type ParquetReaderProps = {
+  defaultDictionarySize?: number;
+};
+
+/** Properties for initializing a ParquetRowGroupReader */
+export type ParquetIterationProps = {
+  /** Filter allowing some columns to be dropped */
+  columnList?: string[] | string[][];
+};
+
+const DEFAULT_PROPS: Required<ParquetReaderProps> = {
+  defaultDictionarySize: 1e6
+};
 
 /**
- * A parquet reader allows retrieving the rows from a parquet file in order.
- * The basic usage is to create a reader and then retrieve a cursor/iterator
- * which allows you to consume row after row until all rows have been read. It is
- * important that you call close() after you are finished reading the file to
- * avoid leaking file descriptors.
+ * The parquet envelope reader allows direct, unbuffered access to the individual
+ * sections of the parquet file, namely the header, footer and the row groups.
+ * This class is intended for advanced/internal users; if you just want to retrieve
+ * rows from a parquet file use the ParquetReader instead
  */
-export class ParquetReader<T> implements AsyncIterable<T> {
-  /**
-   * return a new parquet reader initialized with a read function
-   */
-  static async openBlob<T>(blob: Blob): Promise<ParquetReader<T>> {
-    const readFn = async (start: number, length: number) => {
-      const arrayBuffer = await blob.slice(start, start + length).arrayBuffer();
-      return Buffer.from(arrayBuffer);
-    };
-    const closeFn = async () => {};
-    const size = blob.size;
-    const envelopeReader = new ParquetEnvelopeReader(readFn, closeFn, size);
-    try {
-      await envelopeReader.readHeader();
-      const metadata = await envelopeReader.readFooter();
-      return new ParquetReader(metadata, envelopeReader);
-    } catch (err) {
-      await envelopeReader.close();
-      throw err;
-    }
+export class ParquetReader {
+  props: Required<ParquetReaderProps>;
+  file: ReadableFile;
+  metadata: Promise<FileMetaData> | null = null;
+
+  constructor(file: ReadableFile, props?: ParquetReaderProps) {
+    this.file = file;
+    this.props = {...DEFAULT_PROPS, ...props};
   }
 
-  /**
-   * return a new parquet reader initialized with a read function
-   */
-  static async openArrayBuffer<T>(arrayBuffer: ArrayBuffer): Promise<ParquetReader<T>> {
-    const readFn = async (start: number, length: number) => Buffer.from(arrayBuffer, start, length);
-    const closeFn = async () => {};
-    const size = arrayBuffer.byteLength;
-    const envelopeReader = new ParquetEnvelopeReader(readFn, closeFn, size);
-    try {
-      await envelopeReader.readHeader();
-      const metadata = await envelopeReader.readFooter();
-      return new ParquetReader(metadata, envelopeReader);
-    } catch (err) {
-      await envelopeReader.close();
-      throw err;
+  close(): void {
+    // eslint-disable-next-line @typescript-eslint/no-floating-promises
+    this.file.close();
+  }
+
+  // HIGH LEVEL METHODS
+
+  /** Yield one row at a time */
+  async *rowIterator(props?: ParquetIterationProps) {
+    for await (const rows of this.rowBatchIterator(props)) {
+      // yield *rows
+      for (const row of rows) {
+        yield row;
+      }
     }
   }
 
-  static async openBuffer<T>(buffer: Buffer): Promise<ParquetReader<T>> {
-    const envelopeReader = await ParquetEnvelopeReader.openBuffer(buffer);
-    try {
-      await envelopeReader.readHeader();
-      const metadata = await envelopeReader.readFooter();
-      return new ParquetReader<T>(metadata, envelopeReader);
-    } catch (err) {
-      await envelopeReader.close();
-      throw err;
+  /** Yield one batch of rows at a time */
+  async *rowBatchIterator(props?: ParquetIterationProps) {
+    const schema = await this.getSchema();
+    for await (const rowGroup of this.rowGroupIterator(props)) {
+      yield materializeRecords(schema, rowGroup);
     }
   }
 
-  public metadata: FileMetaData;
-  public envelopeReader: ParquetEnvelopeReader;
-  public schema: ParquetSchema;
+  /** Iterate over the raw row groups */
+  async *rowGroupIterator(props?: ParquetIterationProps) {
+    // Ensure strings are nested in arrays
+    const columnList: string[][] = (props?.columnList || []).map((x) =>
+      Array.isArray(x) ? x : [x]
+    );
 
-  /**
-   * Create a new parquet reader from the file metadata and an envelope reader.
-   * It is not recommended to call this constructor directly except for advanced
-   * and internal use cases. Consider using one of the open{File,Buffer} methods
-   * instead
-   */
-  constructor(metadata: FileMetaData, envelopeReader: ParquetEnvelopeReader) {
-    if (metadata.version !== PARQUET_VERSION) {
-      throw new Error('invalid parquet version');
+    const metadata = await this.getFileMetadata();
+    const schema = await this.getSchema();
+
+    const rowGroupCount = metadata?.row_groups.length || 0;
+
+    for (let rowGroupIndex = 0; rowGroupIndex < rowGroupCount; rowGroupIndex++) {
+      const rowGroup = await this.readRowGroup(
+        schema,
+        metadata.row_groups[rowGroupIndex],
+        columnList
+      );
+      yield rowGroup;
     }
+  }
 
-    this.metadata = metadata;
-    this.envelopeReader = envelopeReader;
-    const root = this.metadata.schema[0];
-    const {schema} = decodeSchema(this.metadata.schema, 1, root.num_children!);
-    this.schema = new ParquetSchema(schema);
+  async getRowCount(): Promise<number> {
+    const metadata = await this.getFileMetadata();
+    return Number(metadata.num_rows);
   }
 
-  /**
-   * Close this parquet reader. You MUST call this method once you're finished
-   * reading rows
-   */
-  async close(): Promise<void> {
-    await this.envelopeReader.close();
-    // this.envelopeReader = null;
-    // this.metadata = null;
+  async getSchema(): Promise<ParquetSchema> {
+    const metadata = await this.getFileMetadata();
+    const root = metadata.schema[0];
+    const {schema: schemaDefinition} = decodeSchema(metadata.schema, 1, root.num_children!);
+    const schema = new ParquetSchema(schemaDefinition);
+    return schema;
   }
 
   /**
-   * Return a cursor to the file. You may open more than one cursor and use
-   * them concurrently. All cursors become invalid once close() is called on
-   * the reader object.
-   *
-   * The required_columns parameter controls which columns are actually read
-   * from disk. An empty array or no value implies all columns. A list of column
-   * names means that only those columns should be loaded from disk.
+   * Returns the user (key/value) metadata for this file
+   * In parquet this is not stored on the schema like it is in arrow
    */
-  getCursor(): ParquetCursor<T>;
-  // @ts-ignore
-  getCursor<K extends keyof T>(columnList: (K | K[])[]): ParquetCursor<Pick<T, K>>;
-  getCursor(columnList: (string | string[])[]): ParquetCursor<Partial<T>>;
-  getCursor(columnList?: (string | string[])[]): ParquetCursor<Partial<T>> {
-    if (!columnList) {
-      // tslint:disable-next-line:no-parameter-reassignment
-      columnList = [];
+  async getSchemaMetadata(): Promise<Record<string, string>> {
+    const metadata = await this.getFileMetadata();
+    const md: Record<string, string> = {};
+    for (const kv of metadata.key_value_metadata!) {
+      md[kv.key] = kv.value!;
     }
+    return md;
+  }
+
+  async getFileMetadata(): Promise<FileMetaData> {
+    if (!this.metadata) {
+      await this.readHeader();
+      this.metadata = this.readFooter();
+    }
+    return this.metadata;
+  }
 
-    // tslint:disable-next-line:no-parameter-reassignment
-    columnList = columnList.map((x) => (Array.isArray(x) ? x : [x]));
+  // LOW LEVEL METHODS
 
-    return new ParquetCursor<T>(
-      this.metadata,
-      this.envelopeReader,
-      this.schema,
-      columnList as string[][]
-    );
+  /** Metadata is stored in the footer */
+  async readHeader(): Promise<void> {
+    const buffer = await this.file.read(0, PARQUET_MAGIC.length);
+    const magic = buffer.toString();
+    switch (magic) {
+      case PARQUET_MAGIC:
+        break;
+      case PARQUET_MAGIC_ENCRYPTED:
+        throw new Error('Encrypted parquet file not supported');
+      default:
+        throw new Error(`Invalid parquet file (magic=${magic})`);
+    }
   }
 
-  /**
-   * Return the number of rows in this file. Note that the number of rows is
-   * not neccessarily equal to the number of rows in each column.
-   */
-  getRowCount(): number {
-    return Number(this.metadata.num_rows);
+  /** Metadata is stored in the footer */
+  async readFooter(): Promise<FileMetaData> {
+    const trailerLen = PARQUET_MAGIC.length + 4;
+    const trailerBuf = await this.file.read(this.file.size - trailerLen, trailerLen);
+
+    const magic = trailerBuf.slice(4).toString();
+    if (magic !== PARQUET_MAGIC) {
+      throw new Error(`Not a valid parquet file (magic="${magic})`);
+    }
+
+    const metadataSize = trailerBuf.readUInt32LE(0);
+    const metadataOffset = this.file.size - metadataSize - trailerLen;
+    if (metadataOffset < PARQUET_MAGIC.length) {
+      throw new Error(`Invalid metadata size ${metadataOffset}`);
+    }
+
+    const metadataBuf = await this.file.read(metadataOffset, metadataSize);
+    // let metadata = new parquet_thrift.FileMetaData();
+    // parquet_util.decodeThrift(metadata, metadataBuf);
+    const {metadata} = decodeFileMetadata(metadataBuf);
+    return metadata;
   }
 
-  /**
-   * Returns the ParquetSchema for this file
-   */
-  getSchema(): ParquetSchema {
-    return this.schema;
+  /** Data is stored in row groups (similar to Apache Arrow record batches) */
+  async readRowGroup(
+    schema: ParquetSchema,
+    rowGroup: RowGroup,
+    columnList: string[][]
+  ): Promise<ParquetBuffer> {
+    const buffer: ParquetBuffer = {
+      rowCount: Number(rowGroup.num_rows),
+      columnData: {}
+    };
+    for (const colChunk of rowGroup.columns) {
+      const colMetadata = colChunk.meta_data;
+      const colKey = colMetadata?.path_in_schema;
+      if (columnList.length > 0 && fieldIndexOf(columnList, colKey!) < 0) {
+        continue; // eslint-disable-line no-continue
+      }
+      buffer.columnData[colKey!.join()] = await this.readColumnChunk(schema, colChunk);
+    }
+    return buffer;
   }
 
   /**
-   * Returns the user (key/value) metadata for this file
+   * Each row group contains column chunks for all the columns.
    */
-  getMetadata(): Record<string, string> {
-    const md: Record<string, string> = {};
-    for (const kv of this.metadata.key_value_metadata!) {
-      md[kv.key] = kv.value!;
+  async readColumnChunk(schema: ParquetSchema, colChunk: ColumnChunk): Promise<ParquetData> {
+    if (colChunk.file_path !== undefined && colChunk.file_path !== null) {
+      throw new Error('external references are not supported');
     }
-    return md;
+
+    const field = schema.findField(colChunk.meta_data?.path_in_schema!);
+    const type: PrimitiveType = getThriftEnum(Type, colChunk.meta_data?.type!) as any;
+
+    if (type !== field.primitiveType) {
+      throw new Error(`chunk type not matching schema: ${type}`);
+    }
+
+    const compression: ParquetCompression = getThriftEnum(
+      CompressionCodec,
+      colChunk.meta_data?.codec!
+    ) as any;
+
+    const pagesOffset = Number(colChunk.meta_data?.data_page_offset!);
+    let pagesSize = Number(colChunk.meta_data?.total_compressed_size!);
+
+    if (!colChunk.file_path) {
+      pagesSize = Math.min(
+        this.file.size - pagesOffset,
+        Number(colChunk.meta_data?.total_compressed_size)
+      );
+    }
+
+    const options: ParquetOptions = {
+      type,
+      rLevelMax: field.rLevelMax,
+      dLevelMax: field.dLevelMax,
+      compression,
+      column: field,
+      numValues: colChunk.meta_data?.num_values,
+      dictionary: []
+    };
+
+    let dictionary;
+
+    const dictionaryPageOffset = colChunk?.meta_data?.dictionary_page_offset;
+
+    if (dictionaryPageOffset) {
+      const dictionaryOffset = Number(dictionaryPageOffset);
+      // Getting dictionary from column chunk to iterate all over indexes to get dataPage values.
+      dictionary = await this.getDictionary(dictionaryOffset, options, pagesOffset);
+    }
+
+    dictionary = options.dictionary?.length ? options.dictionary : dictionary;
+    const pagesBuf = await this.file.read(pagesOffset, pagesSize);
+    return await decodeDataPages(pagesBuf, {...options, dictionary});
   }
 
   /**
-   * Implement AsyncIterable
+   * Getting dictionary for allows to flatten values by indices.
+   * @param dictionaryPageOffset
+   * @param options
+   * @param pagesOffset
+   * @returns
    */
-  // tslint:disable-next-line:function-name
-  [Symbol.asyncIterator](): AsyncIterator<T> {
-    return this.getCursor()[Symbol.asyncIterator]();
+  async getDictionary(
+    dictionaryPageOffset: number,
+    options: ParquetOptions,
+    pagesOffset: number
+  ): Promise<string[]> {
+    if (dictionaryPageOffset === 0) {
+      // dictionarySize = Math.min(this.fileSize - pagesOffset, this.defaultDictionarySize);
+      // pagesBuf = await this.read(pagesOffset, dictionarySize);
+
+      // In this case we are working with parquet-mr files format. Problem is described below:
+      // https://stackoverflow.com/questions/55225108/why-is-dictionary-page-offset-0-for-plain-dictionary-encoding
+      // We need to get dictionary page from column chunk if it exists.
+      // Now if we use code commented above we don't get DICTIONARY_PAGE we get DATA_PAGE instead.
+      return [];
+    }
+
+    const dictionarySize = Math.min(
+      this.file.size - dictionaryPageOffset,
+      this.props.defaultDictionarySize
+    );
+    const pagesBuf = await this.file.read(dictionaryPageOffset, dictionarySize);
+
+    const cursor = {buffer: pagesBuf, offset: 0, size: pagesBuf.length};
+    const decodedPage = await decodePage(cursor, options);
+
+    return decodedPage.dictionary!;
   }
 }