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

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

@@ -0,0 +1,94 @@
+// Forked from https://github.com/kbajalc/parquets under MIT license (Copyright (c) 2017 ironSource Ltd.)
+import {FileMetaData} from '../parquet-thrift';
+import {ParquetEnvelopeReader} from './parquet-envelope-reader';
+import {ParquetSchema} from '../schema/schema';
+import {ParquetRecord} from '../schema/declare';
+import {materializeRecords} from '../schema/shred';
+
+/**
+ * A parquet cursor is used to retrieve rows from a parquet file in order
+ */
+export class ParquetCursor<T> implements AsyncIterable<T> {
+  public metadata: FileMetaData;
+  public envelopeReader: ParquetEnvelopeReader;
+  public schema: ParquetSchema;
+  public columnList: string[][];
+  public rowGroup: ParquetRecord[] = [];
+  public rowGroupIndex: number;
+
+  /**
+   * Create a new parquet reader from the file metadata and an envelope reader.
+   * It is usually not recommended to call this constructor directly except for
+   * advanced and internal use cases. Consider using getCursor() on the
+   * ParquetReader instead
+   */
+  constructor(
+    metadata: FileMetaData,
+    envelopeReader: ParquetEnvelopeReader,
+    schema: ParquetSchema,
+    columnList: string[][]
+  ) {
+    this.metadata = metadata;
+    this.envelopeReader = envelopeReader;
+    this.schema = schema;
+    this.columnList = columnList;
+    this.rowGroupIndex = 0;
+  }
+
+  /**
+   * Retrieve the next row from the cursor. Returns a row or NULL if the end
+   * of the file was reached
+   */
+  async next<T = any>(): Promise<T> {
+    if (this.rowGroup.length === 0) {
+      if (this.rowGroupIndex >= this.metadata.row_groups.length) {
+        // @ts-ignore
+        return null;
+      }
+      const rowBuffer = await this.envelopeReader.readRowGroup(
+        this.schema,
+        this.metadata.row_groups[this.rowGroupIndex],
+        this.columnList
+      );
+      this.rowGroup = materializeRecords(this.schema, rowBuffer);
+      this.rowGroupIndex++;
+    }
+    return this.rowGroup.shift() as any;
+  }
+
+  /**
+   * Rewind the cursor the the beginning of the file
+   */
+  rewind(): void {
+    this.rowGroup = [];
+    this.rowGroupIndex = 0;
+  }
+
+  /**
+   * Implement AsyncIterable
+   */
+  // tslint:disable-next-line:function-name
+  [Symbol.asyncIterator](): AsyncIterator<T> {
+    let done = false;
+    return {
+      next: async () => {
+        if (done) {
+          return {done, value: null};
+        }
+        const value = await this.next();
+        if (value === null) {
+          return {done: true, value};
+        }
+        return {done: false, value};
+      },
+      return: async () => {
+        done = true;
+        return {done, value: null};
+      },
+      throw: async () => {
+        done = true;
+        return {done: true, value: null};
+      }
+    };
+  }
+}

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

@@ -0,0 +1,210 @@
+// Forked from https://github.com/kbajalc/parquets under MIT license (Copyright (c) 2017 ironSource Ltd.)
+import {ParquetSchema} from '../schema/schema';
+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 {fstat, fopen, fread, fclose} from '../utils/file-utils';
+import {decodeFileMetadata, getThriftEnum, fieldIndexOf} from '../utils/read-utils';
+import {decodeDataPages, decodePage} from './decoders';
+
+const DEFAULT_DICTIONARY_SIZE = 1e6;
+
+/**
+ * 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 ParquetEnvelopeReader {
+  public read: (position: number, length: number) => Promise<Buffer>;
+  /**
+   * Close this parquet reader. You MUST call this method once you're finished
+   * reading rows
+   */
+  public close: () => Promise<void>;
+  public fileSize: number;
+  public defaultDictionarySize: number;
+
+  static async openFile(filePath: string): Promise<ParquetEnvelopeReader> {
+    const fileStat = await fstat(filePath);
+    const fileDescriptor = await fopen(filePath);
+
+    const readFn = fread.bind(undefined, fileDescriptor);
+    const closeFn = fclose.bind(undefined, fileDescriptor);
+
+    return new ParquetEnvelopeReader(readFn, closeFn, fileStat.size);
+  }
+
+  static async openBuffer(buffer: Buffer): Promise<ParquetEnvelopeReader> {
+    const readFn = (position: number, length: number) =>
+      Promise.resolve(buffer.slice(position, position + length));
+    const closeFn = () => Promise.resolve();
+    return new ParquetEnvelopeReader(readFn, closeFn, buffer.length);
+  }
+
+  constructor(
+    read: (position: number, length: number) => Promise<Buffer>,
+    close: () => Promise<void>,
+    fileSize: number,
+    options?: any
+  ) {
+    this.read = read;
+    this.close = close;
+    this.fileSize = fileSize;
+    this.defaultDictionarySize = options?.defaultDictionarySize || DEFAULT_DICTIONARY_SIZE;
+  }
+
+  async readHeader(): Promise<void> {
+    const buffer = await this.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})`);
+    }
+  }
+
+  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;
+  }
+
+  /**
+   * Do reading of parquet file's column chunk
+   * @param schema
+   * @param colChunk
+   */
+  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');
+    }
+
+    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.fileSize - 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.read(pagesOffset, pagesSize);
+    return await decodeDataPages(pagesBuf, {...options, dictionary});
+  }
+
+  /**
+   * Getting dictionary for allows to flatten values by indices.
+   * @param dictionaryPageOffset
+   * @param options
+   * @param pagesOffset
+   * @returns
+   */
+  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.fileSize - dictionaryPageOffset,
+      this.defaultDictionarySize
+    );
+    const pagesBuf = await this.read(dictionaryPageOffset, dictionarySize);
+
+    const cursor = {buffer: pagesBuf, offset: 0, size: pagesBuf.length};
+    const decodedPage = await decodePage(cursor, options);
+
+    return decodedPage.dictionary!;
+  }
+
+  async readFooter(): Promise<FileMetaData> {
+    const trailerLen = PARQUET_MAGIC.length + 4;
+    const trailerBuf = await this.read(this.fileSize - 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.fileSize - metadataSize - trailerLen;
+    if (metadataOffset < PARQUET_MAGIC.length) {
+      throw new Error(`Invalid metadata size ${metadataOffset}`);
+    }
+
+    const metadataBuf = await this.read(metadataOffset, metadataSize);
+    // let metadata = new parquet_thrift.FileMetaData();
+    // parquet_util.decodeThrift(metadata, metadataBuf);
+    const {metadata} = decodeFileMetadata(metadataBuf);
+    return metadata;
+  }
+}

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

@@ -0,0 +1,179 @@
+// 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 {ParquetSchema} from '../schema/schema';
+import {ParquetCursor} from './parquet-cursor';
+import {PARQUET_VERSION} from '../../constants';
+import {decodeSchema} from './decoders';
+
+/**
+ * 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.
+ */
+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;
+    }
+  }
+
+  /**
+   * 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;
+    }
+  }
+
+  /**
+   * Open the parquet file pointed to by the specified path and return a new
+   * parquet reader
+   */
+  static async openFile<T>(filePath: string): Promise<ParquetReader<T>> {
+    const envelopeReader = await ParquetEnvelopeReader.openFile(filePath);
+    try {
+      await envelopeReader.readHeader();
+      const metadata = await envelopeReader.readFooter();
+      return new ParquetReader<T>(metadata, envelopeReader);
+    } catch (err) {
+      await envelopeReader.close();
+      throw err;
+    }
+  }
+
+  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;
+    }
+  }
+
+  public metadata: FileMetaData;
+  public envelopeReader: ParquetEnvelopeReader;
+  public schema: ParquetSchema;
+
+  /**
+   * 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');
+    }
+
+    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);
+  }
+
+  /**
+   * 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;
+  }
+
+  /**
+   * 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.
+   */
+  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 = [];
+    }
+
+    // tslint:disable-next-line:no-parameter-reassignment
+    columnList = columnList.map((x) => (Array.isArray(x) ? x : [x]));
+
+    return new ParquetCursor<T>(
+      this.metadata,
+      this.envelopeReader,
+      this.schema,
+      columnList as string[][]
+    );
+  }
+
+  /**
+   * 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);
+  }
+
+  /**
+   * Returns the ParquetSchema for this file
+   */
+  getSchema(): ParquetSchema {
+    return this.schema;
+  }
+
+  /**
+   * Returns the user (key/value) metadata for this file
+   */
+  getMetadata(): Record<string, string> {
+    const md: Record<string, string> = {};
+    for (const kv of this.metadata.key_value_metadata!) {
+      md[kv.key] = kv.value!;
+    }
+    return md;
+  }
+
+  /**
+   * Implement AsyncIterable
+   */
+  // tslint:disable-next-line:function-name
+  [Symbol.asyncIterator](): AsyncIterator<T> {
+    return this.getCursor()[Symbol.asyncIterator]();
+  }
+}