@loaders.gl/parquet 4.0.0-alpha.5 → 4.0.0-alpha.7

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 {materializeRows} from '../schema/shred';
+
+import {PARQUET_MAGIC, PARQUET_MAGIC_ENCRYPTED} from '../../constants';
+import {ColumnChunk, CompressionCodec, FileMetaData, RowGroup, Type} from '../parquet-thrift';
+import {
+  ParquetRowGroup,
+  ParquetCompression,
+  ParquetColumnChunk,
+  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 materializeRows(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<ParquetRowGroup> {
+    const buffer: ParquetRowGroup = {
+      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<ParquetColumnChunk> {
+    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!;
   }
 }

package/src/parquetjs/schema/declare.ts

@@ -98,6 +98,7 @@ export interface ParquetField {
   fields?: Record<string, ParquetField>;
 }
 
+/** @todo better name, this is an internal type? */
 export interface ParquetOptions {
   type: ParquetType;
   rLevelMax: number;
@@ -108,32 +109,40 @@ export interface ParquetOptions {
   dictionary?: ParquetDictionary;
 }
 
-export interface ParquetData {
-  dlevels: number[];
-  rlevels: number[];
-  values: any[];
-  count: number;
-  pageHeaders: PageHeader[];
-}
-
 export interface ParquetPageData {
   dlevels: number[];
   rlevels: number[];
-  values: any[];
+  /** Actual column chunks */
+  values: any[]; // ArrayLike<any>;
   count: number;
   dictionary?: ParquetDictionary;
+  /** The "raw" page header from the file */
   pageHeader: PageHeader;
 }
 
-export interface ParquetRecord {
+export interface ParquetRow {
   [key: string]: any;
 }
 
-export class ParquetBuffer {
+/** @
+ * Holds data for one row group (column chunks) */
+export class ParquetRowGroup {
+  /** Number of rows in this page */
   rowCount: number;
-  columnData: Record<string, ParquetData>;
-  constructor(rowCount: number = 0, columnData: Record<string, ParquetData> = {}) {
+  /** Map of Column chunks */
+  columnData: Record<string, ParquetColumnChunk>;
+
+  constructor(rowCount: number = 0, columnData: Record<string, ParquetColumnChunk> = {}) {
     this.rowCount = rowCount;
     this.columnData = columnData;
   }
 }
+
+/** Holds the data for one column chunk */
+export interface ParquetColumnChunk {
+  dlevels: number[];
+  rlevels: number[];
+  values: any[];
+  count: number;
+  pageHeaders: PageHeader[];
+}

package/src/parquetjs/schema/schema.ts

@@ -4,14 +4,14 @@ import {PARQUET_CODECS} from '../codecs';
 import {PARQUET_COMPRESSION_METHODS} from '../compression';
 import {
   FieldDefinition,
-  ParquetBuffer,
+  ParquetRowGroup,
   ParquetCompression,
   ParquetField,
-  ParquetRecord,
+  ParquetRow,
   RepetitionType,
   SchemaDefinition
 } from './declare';
-import {materializeRecords, shredBuffer, shredRecord} from './shred';
+import {materializeRows, shredBuffer, shredRecord} from './shred';
 import {PARQUET_LOGICAL_TYPES} from './types';
 
 /**
@@ -70,12 +70,12 @@ export class ParquetSchema {
     return branch;
   }
 
-  shredRecord(record: ParquetRecord, buffer: ParquetBuffer): void {
-    shredRecord(this, record, buffer);
+  shredRecord(row: ParquetRow, rowGroup: ParquetRowGroup): void {
+    shredRecord(this, row, rowGroup);
   }
 
-  materializeRecords(buffer: ParquetBuffer): ParquetRecord[] {
-    return materializeRecords(this, buffer);
+  materializeRows(rowGroup: ParquetRowGroup): ParquetRow[] {
+    return materializeRows(this, rowGroup);
   }
 
   compress(type: ParquetCompression): this {
@@ -84,7 +84,7 @@ export class ParquetSchema {
     return this;
   }
 
-  buffer(): ParquetBuffer {
+  rowGroup(): ParquetRowGroup {
     return shredBuffer(this);
   }
 }