@loaders.gl/parquet 4.4.5 → 5.0.0-alpha.2

package/src/lib/parsers/parse-parquet-to-arrow.ts

@@ -5,68 +5,95 @@
 // eslint-disable
 import type {ReadableFile} from '@loaders.gl/loader-utils';
 import type {ArrowTable, ArrowTableBatch, Schema} from '@loaders.gl/schema';
-import {convertArrowToSchema} from '@loaders.gl/schema-utils';
 
-import type * as parquetWasm from 'parquet-wasm';
+import type * as parquetWasm from 'parquet-wasm/esm/parquet_wasm.js';
 import * as arrow from 'apache-arrow';
 
+import {
+  normalizeArrowTableGeoMetadata,
+  getParquetFileMetadataMap
+} from '../geo/geospatial-metadata';
 import {loadWasm} from '../utils/load-wasm';
 import {makeStreamIterator} from '../utils/make-stream-iterator';
 
+type ParquetWasmReaderOptions = parquetWasm.ReaderOptions & {wasmUrl?: string};
+
 export async function parseParquetFileToArrow(
   file: ReadableFile,
-  options?: parquetWasm.ReaderOptions & {wasmUrl: string}
+  options?: ParquetWasmReaderOptions
 ): Promise<ArrowTable> {
   const wasmUrl = options?.wasmUrl;
+  const readerOptions = getReaderOptions(options);
 
   const wasm = await loadWasm(wasmUrl);
 
-  let parquetFile: parquetWasm.ParquetFile;
-  if (file.handle instanceof Blob) {
-    // TODO - let's assume fromFile() works on Blobs and not just on File...
-    parquetFile = await wasm.ParquetFile.fromFile(file.handle as File);
-  } else {
-    parquetFile = await wasm.ParquetFile.fromUrl(file.url);
+  const parquetFile = await createParquetFile(file, wasm);
+  try {
+    const wasmTable = await parquetFile.read(readerOptions);
+    const ipcStream = wasmTable.intoIPCStream();
+    const arrowTable = arrow.tableFromIPC(ipcStream);
+    return normalizeArrowTableGeoMetadata(
+      {
+        shape: 'arrow-table',
+        data: arrowTable
+      },
+      getParquetFileMetadataMap(parquetFile.metadata())
+    );
+  } finally {
+    parquetFile.free();
   }
-
-  const wasmTable = await parquetFile.read(options);
-  const ipcStream = wasmTable.intoIPCStream();
-  const arrowTable = arrow.tableFromIPC(ipcStream);
-
-  return {
-    shape: 'arrow-table',
-    schema: convertArrowToSchema(arrowTable.schema),
-    data: arrowTable
-  };
 }
 
 export async function* parseParquetFileToArrowInBatches(
   file: ReadableFile,
-  options: parquetWasm.ReaderOptions & {wasmUrl: string}
+  options?: ParquetWasmReaderOptions
 ): AsyncIterable<ArrowTableBatch> {
   const wasmUrl = options?.wasmUrl;
+  const readerOptions = getReaderOptions(options);
 
   const wasm = await loadWasm(wasmUrl);
 
-  let parquetFile: parquetWasm.ParquetFile;
+  const parquetFile = await createParquetFile(file, wasm);
+  try {
+    const stream = await parquetFile.stream(readerOptions);
+    const schemaMetadata = getParquetFileMetadataMap(parquetFile.metadata());
+
+    let schema: Schema | undefined;
+    for await (const wasmRecordBatch of makeStreamIterator<parquetWasm.RecordBatch>(stream)) {
+      const normalizedArrowTable = normalizeArrowTableGeoMetadata(
+        {
+          shape: 'arrow-table',
+          data: arrow.tableFromIPC(wasmRecordBatch.intoIPCStream())
+        },
+        schemaMetadata
+      );
+      schema ||= normalizedArrowTable.schema;
+      yield {
+        batchType: 'data',
+        shape: 'arrow-table',
+        schema,
+        data: normalizedArrowTable.data,
+        length: normalizedArrowTable.data.numRows
+      };
+    }
+  } finally {
+    parquetFile.free();
+  }
+}
+
+async function createParquetFile(
+  file: ReadableFile,
+  wasm: typeof parquetWasm
+): Promise<parquetWasm.ParquetFile> {
   if (file.handle instanceof Blob) {
     // Works on Blobs: https://kylebarron.dev/parquet-wasm/classes/esm_parquet_wasm.ParquetFile.html#fromFile
-    parquetFile = await wasm.ParquetFile.fromFile(file.handle as File);
-  } else {
-    parquetFile = await wasm.ParquetFile.fromUrl(file.url);
+    return await wasm.ParquetFile.fromFile(file.handle as File);
   }
 
-  const stream: ReadableStream<arrow.RecordBatch> = await parquetFile.stream(options);
+  return await wasm.ParquetFile.fromUrl(file.url);
+}
 
-  let schema: Schema;
-  for await (const recordBatch of makeStreamIterator(stream)) {
-    schema ||= convertArrowToSchema(recordBatch.schema);
-    yield {
-      batchType: 'data',
-      shape: 'arrow-table',
-      schema,
-      data: new arrow.Table(recordBatch),
-      length: recordBatch.numRows
-    };
-  }
+function getReaderOptions(options?: ParquetWasmReaderOptions): parquetWasm.ReaderOptions {
+  const {wasmUrl, ...readerOptions} = options || {};
+  return readerOptions;
 }

package/src/lib/parsers/parse-parquet-to-columns.ts

@@ -4,13 +4,12 @@
 
 import type {ColumnarTable, ColumnarTableBatch, Schema} from '@loaders.gl/schema';
 import type {ReadableFile} from '@loaders.gl/loader-utils';
-import type {ParquetJSONLoaderOptions} from '../../parquet-json-loader';
+import type {ParquetLoaderOptions} from '../../parquet-loader';
 import {ParquetReader} from '../../parquetjs/parser/parquet-reader';
 import {ParquetRowGroup} from '../../parquetjs/schema/declare';
 import {ParquetSchema} from '../../parquetjs/schema/schema';
 import {materializeColumns} from '../../parquetjs/schema/shred';
 import {getSchemaFromParquetReader} from './get-parquet-schema';
-import {installBufferPolyfill} from '../../polyfills/buffer/index';
 import {preloadCompressions} from '../../parquetjs/compression';
 
 /**
@@ -18,9 +17,8 @@ import {preloadCompressions} from '../../parquetjs/compression';
  */
 export async function parseParquetFileInColumns(
   file: ReadableFile,
-  options?: ParquetJSONLoaderOptions
+  options?: ParquetLoaderOptions
 ): Promise<ColumnarTable> {
-  installBufferPolyfill();
   await preloadCompressions(options);
 
   for await (const batch of parseParquetFileInColumnarBatches(file, options)) {
@@ -38,9 +36,8 @@ export async function parseParquetFileInColumns(
  */
 export async function* parseParquetFileInColumnarBatches(
   file: ReadableFile,
-  options?: ParquetJSONLoaderOptions
+  options?: ParquetLoaderOptions
 ): AsyncIterable<ColumnarTableBatch> {
-  installBufferPolyfill();
   await preloadCompressions(options);
 
   const reader = new ParquetReader(file);
@@ -51,7 +48,9 @@ export async function* parseParquetFileInColumnarBatches(
   const parquetSchema = await reader.getSchema();
 
   // Iterate over row batches
-  const rowGroups = reader.rowGroupIterator(options?.parquet);
+  const rowGroups = reader.rowGroupIterator(
+    options?.parquet?.columns?.length ? {columnList: options.parquet.columns} : undefined
+  );
   for await (const rowGroup of rowGroups) {
     yield convertRowGroupToTableBatch(rowGroup, parquetSchema, schema);
   }

package/src/lib/parsers/parse-parquet-to-json.ts

@@ -4,13 +4,12 @@
 
 import {default as log} from '@probe.gl/log';
 import type {ReadableFile} from '@loaders.gl/loader-utils';
-import type {ObjectRowTable, ObjectRowTableBatch} from '@loaders.gl/schema';
+import type {ObjectRowTable, ObjectRowTableBatch, Schema} from '@loaders.gl/schema';
 
-import type {ParquetJSONLoaderOptions} from '../../parquet-json-loader';
+import type {ParquetJSLoaderOptions} from '../../parquet-loader-options';
 import type {ParquetRow} from '../../parquetjs/schema/declare';
 import {ParquetReader} from '../../parquetjs/parser/parquet-reader';
 import {getSchemaFromParquetReader} from './get-parquet-schema';
-import {installBufferPolyfill} from '../../polyfills/buffer/index';
 import {preloadCompressions} from '../../parquetjs/compression';
 
 /**
@@ -21,9 +20,8 @@ import {preloadCompressions} from '../../parquetjs/compression';
  */
 export async function parseParquetFile(
   file: ReadableFile,
-  options?: ParquetJSONLoaderOptions
+  options?: ParquetJSLoaderOptions
 ): Promise<ObjectRowTable> {
-  installBufferPolyfill();
   await preloadCompressions(options);
 
   const reader = new ParquetReader(file, {
@@ -31,30 +29,37 @@ export async function parseParquetFile(
   });
 
   const schema = await getSchemaFromParquetReader(reader);
+  const projectedSchema = projectSchema(schema, options?.parquet?.columns);
 
   const rows: ParquetRow[] = [];
 
-  const rowBatches = reader.rowBatchIterator(options?.parquet);
+  const rowBatches = reader.rowBatchIterator(getParquetIterationProps(options));
+  const rowOffset = options?.parquet?.offset || 0;
+  const rowLimit = options?.parquet?.limit;
+  let rowsSkipped = 0;
   for await (const rowBatch of rowBatches) {
-    let limitHasReached = false;
-    // we have only one input batch so return
     for (const row of rowBatch) {
-      if (options?.limit && rows.length >= options?.core?.limit!) {
-        limitHasReached = true;
-        break;
+      if (rowsSkipped < rowOffset) {
+        rowsSkipped++;
+        continue;
+      }
+
+      if (rowLimit !== undefined && rows.length >= rowLimit) {
+        log.warn(`Rows number limit has been reached. Only first ${rowLimit} are loaded`)();
+        return {
+          shape: 'object-row-table',
+          schema: projectedSchema,
+          data: rows
+        };
       }
+
       rows.push(row);
     }
-    if (limitHasReached) {
-      log.warn(
-        `Rows number limit has been reached. Only first ${options?.core?.limit} are loaded`
-      )();
-      break;
-    }
   }
+
   const objectRowTable: ObjectRowTable = {
     shape: 'object-row-table',
-    schema,
+    schema: projectedSchema,
     data: rows
   };
 
@@ -68,9 +73,8 @@ export async function parseParquetFile(
  */
 export async function* parseParquetFileInBatches(
   file: ReadableFile,
-  options?: ParquetJSONLoaderOptions
+  options?: ParquetJSLoaderOptions
 ): AsyncIterable<ObjectRowTableBatch> {
-  installBufferPolyfill();
   await preloadCompressions(options);
 
   const reader = new ParquetReader(file, {
@@ -78,19 +82,83 @@ export async function* parseParquetFileInBatches(
   });
 
   const schema = await getSchemaFromParquetReader(reader);
-  const rowBatches = reader.rowBatchIterator(options?.parquet);
+  const projectedSchema = projectSchema(schema, options?.parquet?.columns);
+  const rowBatches = reader.rowBatchIterator(getParquetIterationProps(options));
+  const rowOffset = options?.parquet?.offset || 0;
+  const rowLimit = options?.parquet?.limit;
+  const outputBatchSize = options?.parquet?.batchSize;
+  let rowsSkipped = 0;
+  let rowsYielded = 0;
+  let reachedLimit = false;
+
   for await (const rows of rowBatches) {
-    const objectRowTable: ObjectRowTable = {
-      shape: 'object-row-table',
-      schema,
-      data: rows
-    };
-
-    yield {
-      batchType: 'data',
-      schema,
-      ...objectRowTable,
-      length: rows.length
-    };
+    const nextRows: ParquetRow[] = [];
+    for (const row of rows) {
+      if (rowsSkipped < rowOffset) {
+        rowsSkipped++;
+        continue;
+      }
+
+      if (rowLimit !== undefined && rowsYielded >= rowLimit) {
+        reachedLimit = true;
+        break;
+      }
+
+      nextRows.push(row);
+      rowsYielded++;
+    }
+
+    if (nextRows.length === 0) {
+      continue;
+    }
+
+    if (!outputBatchSize || outputBatchSize <= 0) {
+      yield makeObjectRowBatch(projectedSchema, nextRows);
+      continue;
+    }
+
+    for (let rowIndex = 0; rowIndex < nextRows.length; rowIndex += outputBatchSize) {
+      yield makeObjectRowBatch(
+        projectedSchema,
+        nextRows.slice(rowIndex, rowIndex + outputBatchSize)
+      );
+    }
+
+    if (reachedLimit) {
+      return;
+    }
   }
 }
+
+function getParquetIterationProps(
+  options?: ParquetJSLoaderOptions
+): {columnList?: string[] | string[][]} | undefined {
+  const columnList = options?.parquet?.columns;
+  return columnList?.length ? {columnList} : undefined;
+}
+
+function makeObjectRowBatch(schema: Schema, rows: ParquetRow[]): ObjectRowTableBatch {
+  const objectRowTable: ObjectRowTable = {
+    shape: 'object-row-table',
+    schema,
+    data: rows
+  };
+
+  return {
+    batchType: 'data',
+    schema,
+    ...objectRowTable,
+    length: rows.length
+  };
+}
+
+function projectSchema(schema: Schema, columns?: string[]): Schema {
+  if (!columns?.length) {
+    return schema;
+  }
+
+  return {
+    ...schema,
+    fields: schema.fields.filter(field => columns.includes(field.name))
+  };
+}