@loaders.gl/parquet 3.3.0 → 3.4.0-alpha.1

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

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

package/src/parquetjs/file.ts

@@ -0,0 +1,90 @@
+// Forked from https://github.com/kbajalc/parquets under MIT license (Copyright (c) 2017 ironSource Ltd.)
+import fs from 'fs';
+
+export function fopen(filePath) {
+  return new Promise((resolve, reject) => {
+    fs.open(filePath, 'r', (err, fd) => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve(fd);
+      }
+    });
+  });
+}
+
+export function fstat(filePath) {
+  return new Promise<fs.Stats>((resolve, reject) => {
+    fs.stat(filePath, (err, stat) => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve(stat);
+      }
+    });
+  });
+}
+
+export function fread(fd, position, length) {
+  const buffer = Buffer.alloc(length);
+
+  return new Promise((resolve, reject) => {
+    fs.read(fd, buffer, 0, length, position, (err, bytesRead, buf) => {
+      if (err || bytesRead !== length) {
+        reject(err || Error('read failed'));
+      } else {
+        resolve(buf);
+      }
+    });
+  });
+}
+
+export function fclose(fd) {
+  return new Promise((resolve, reject) => {
+    fs.close(fd, (err) => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve(err);
+      }
+    });
+  });
+}
+
+export function oswrite(os, buf): Promise<void> {
+  return new Promise((resolve, reject) => {
+    os.write(buf, (err) => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve();
+      }
+    });
+  });
+}
+
+export function osclose(os): Promise<void> {
+  return new Promise((resolve, reject) => {
+    os.close((err) => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve();
+      }
+    });
+  });
+}
+
+export function osopen(path, opts) {
+  return new Promise((resolve, reject) => {
+    const outputStream = fs.createWriteStream(path, opts);
+
+    outputStream.on('open', function (fd) {
+      resolve(outputStream);
+    });
+
+    outputStream.on('error', function (err) {
+      reject(err);
+    });
+  });
+}

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,199 @@
+// 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 {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 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;
+  }
+}