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

package/src/parquetjs/schema/schema.ts

@@ -5,6 +5,8 @@
 // Forked from https://github.com/kbajalc/parquets under MIT license
 
 import {PARQUET_CODECS} from '../codecs/index';
+import {isByteStreamSplitType} from '../codecs/byte-stream-split';
+import {isDeltaEncodingType} from '../codecs/delta';
 import {PARQUET_COMPRESSION_METHODS} from '../compression';
 import {
   FieldDefinition,
@@ -15,7 +17,8 @@ import {
   RepetitionType,
   SchemaDefinition
 } from './declare';
-import {materializeRows, shredBuffer, shredRecord} from './shred';
+import type {ArrayType} from '@loaders.gl/schema';
+import {materializeColumns, materializeRows, shredBuffer, shredRecord} from './shred';
 import {PARQUET_LOGICAL_TYPES} from './types';
 
 /**
@@ -82,6 +85,11 @@ export class ParquetSchema {
     return materializeRows(this, rowGroup);
   }
 
+  /** Materializes one decoded row group as top-level column arrays. */
+  materializeColumns(rowGroup: ParquetRowGroup): Record<string, ArrayType> {
+    return materializeColumns(this, rowGroup);
+  }
+
   compress(type: ParquetCompression): this {
     setCompress(this.schema, type);
     setCompress(this.fields, type);
@@ -140,6 +148,8 @@ function buildFields(
         name,
         path: cpath,
         key: cpath.join(),
+        logicalType: opts.logicalType,
+        fieldId: opts.fieldId,
         repetitionType,
         rLevelMax,
         dLevelMax,
@@ -159,6 +169,13 @@ function buildFields(
     if (!(opts.encoding in PARQUET_CODECS)) {
       throw new Error(`unsupported parquet encoding: ${opts.encoding}`);
     }
+    const primitiveType = opts.physicalType || typeDef.primitiveType;
+    if (opts.encoding === 'BYTE_STREAM_SPLIT' && !isByteStreamSplitType(primitiveType)) {
+      throw new Error(`BYTE_STREAM_SPLIT does not support ${primitiveType}`);
+    }
+    if (opts.encoding.startsWith('DELTA_') && !isDeltaEncodingType(opts.encoding, primitiveType)) {
+      throw new Error(`${opts.encoding} does not support ${primitiveType}`);
+    }
 
     opts.compression = opts.compression || 'UNCOMPRESSED';
     if (!(opts.compression in PARQUET_COMPRESSION_METHODS)) {
@@ -169,15 +186,18 @@ function buildFields(
     const cpath = path.concat([name]);
     fieldList[name] = {
       name,
-      primitiveType: typeDef.primitiveType,
+      primitiveType,
       originalType: typeDef.originalType,
+      logicalType: opts.logicalType,
+      fieldId: opts.fieldId,
       path: cpath,
       key: cpath.join(),
       repetitionType,
       encoding: opts.encoding,
       compression: opts.compression,
       typeLength: opts.typeLength || typeDef.typeLength,
-      presision: opts.presision,
+      presision: opts.precision ?? opts.presision,
+      precision: opts.precision ?? opts.presision,
       scale: opts.scale,
       rLevelMax,
       dLevelMax

package/src/parquetjs/schema/shred.ts

@@ -65,9 +65,18 @@ export function shredRecord(
   }
   rowGroup.rowCount += 1;
   for (const field of schema.fieldList) {
-    Array.prototype.push.apply(rowGroup.columnData[field.key].rlevels, data[field.key].rlevels);
-    Array.prototype.push.apply(rowGroup.columnData[field.key].dlevels, data[field.key].dlevels);
-    Array.prototype.push.apply(rowGroup.columnData[field.key].values, data[field.key].values);
+    Array.prototype.push.apply(
+      rowGroup.columnData[field.key].rlevels as number[],
+      data[field.key].rlevels as number[]
+    );
+    Array.prototype.push.apply(
+      rowGroup.columnData[field.key].dlevels as number[],
+      data[field.key].dlevels as number[]
+    );
+    Array.prototype.push.apply(
+      rowGroup.columnData[field.key].values as unknown[],
+      data[field.key].values as unknown as unknown[]
+    );
     rowGroup.columnData[field.key].count += data[field.key].count;
   }
 }
@@ -111,8 +120,8 @@ function shredRecordFields(
         shredRecordFields(field.fields!, null!, data, rLevel, dLevel);
       } else {
         data[field.key].count += 1;
-        data[field.key].rlevels.push(rLevel);
-        data[field.key].dlevels.push(dLevel);
+        (data[field.key].rlevels as number[]).push(rLevel);
+        (data[field.key].dlevels as number[]).push(dLevel);
       }
       continue; // eslint-disable-line no-continue
     }
@@ -124,10 +133,10 @@ function shredRecordFields(
         shredRecordFields(field.fields!, values[i], data, rlvl, field.dLevelMax);
       } else {
         data[field.key].count += 1;
-        data[field.key].rlevels.push(rlvl);
-        data[field.key].dlevels.push(field.dLevelMax);
-        data[field.key].values.push(
-          Types.toPrimitive((field.originalType || field.primitiveType)!, values[i])
+        (data[field.key].rlevels as number[]).push(rlvl);
+        (data[field.key].dlevels as number[]).push(field.dLevelMax);
+        (data[field.key].values as unknown[]).push(
+          Types.toPrimitive((field.originalType || field.primitiveType)!, values[i], field)
         );
       }
     }
@@ -154,10 +163,9 @@ function shredRecordFields(
  *   }
  */
 export function materializeRows(schema: ParquetSchema, rowGroup: ParquetRowGroup): ParquetRow[] {
-  const rows: ParquetRow[] = [];
-  // rows = new Array(rowGroup.rowCount).fill({})'
+  const rows = new Array<ParquetRow>(rowGroup.rowCount);
   for (let i = 0; i < rowGroup.rowCount; i++) {
-    rows.push({});
+    rows[i] = {};
   }
   for (const key in rowGroup.columnData) {
     const columnData = rowGroup.columnData[key];
@@ -179,6 +187,14 @@ function materializeColumnAsRows(
   const field = schema.findField(key);
   const branch = schema.findFieldBranch(key);
 
+  if (branch.length === 1 && field.repetitionType !== 'REPEATED') {
+    materializeFlatColumnAsRows(field, columnData, rows);
+    return;
+  }
+
+  const logicalType = field.originalType || field.primitiveType!;
+  const fromPrimitive = Types.PARQUET_LOGICAL_TYPES[logicalType].fromPrimitive;
+
   // tslint:disable-next-line:prefer-array-literal
   const rLevels: number[] = new Array(field.rLevelMax + 1).fill(0);
   let vIndex = 0;
@@ -219,12 +235,8 @@ function materializeColumnAsRows(
 
     // Leaf node - Add the value
     if (dLevel === field.dLevelMax) {
-      const value = Types.fromPrimitive(
-        // @ts-ignore
-        field.originalType || field.primitiveType,
-        columnData.values[vIndex],
-        field
-      );
+      const primitiveValue = columnData.values[vIndex];
+      const value = fromPrimitive ? fromPrimitive(primitiveValue, field) : primitiveValue;
       vIndex++;
 
       switch (field.repetitionType) {
@@ -248,6 +260,35 @@ function materializeColumnAsRows(
   }
 }
 
+/** Materializes a top-level required or optional primitive directly into row objects. */
+function materializeFlatColumnAsRows(
+  field: ParquetField,
+  columnData: ParquetColumnChunk,
+  rows: ParquetRow[]
+): void {
+  const logicalType = field.originalType || field.primitiveType!;
+  const fromPrimitive = Types.PARQUET_LOGICAL_TYPES[logicalType].fromPrimitive;
+  const count = Math.min(columnData.count, rows.length);
+  let valueIndex = 0;
+
+  if (field.repetitionType === 'REQUIRED' && !fromPrimitive) {
+    for (let rowIndex = 0; rowIndex < count; rowIndex++) {
+      rows[rowIndex][field.name] = columnData.values[rowIndex];
+    }
+    return;
+  }
+
+  for (let rowIndex = 0; rowIndex < count; rowIndex++) {
+    if (columnData.dlevels[rowIndex] === field.dLevelMax) {
+      const primitiveValue = columnData.values[valueIndex];
+      rows[rowIndex][field.name] = fromPrimitive
+        ? fromPrimitive(primitiveValue, field)
+        : primitiveValue;
+      valueIndex++;
+    }
+  }
+}
+
 // Columnar export
 
 /**
@@ -275,14 +316,30 @@ export function materializeColumns(
 ): Record<string, ArrayType> {
   const columns: Record<string, ArrayType> = {};
   for (const key in rowGroup.columnData) {
-    const columnData = rowGroup.columnData[key];
-    if (columnData.count) {
-      materializeColumnAsColumnarArray(schema, columnData, rowGroup.rowCount, key, columns);
+    const column = materializeColumn(schema, rowGroup, key);
+    if (column) {
+      columns[schema.findFieldBranch(key)[0].name] = column;
     }
   }
   return columns;
 }
 
+/** Materializes one decoded Parquet column into its top-level columnar representation. */
+export function materializeColumn(
+  schema: ParquetSchema,
+  rowGroup: ParquetRowGroup,
+  key: string
+): ArrayType | undefined {
+  const columnData = rowGroup.columnData[key];
+  if (!columnData?.count) {
+    return undefined;
+  }
+
+  const columns: Record<string, ArrayType> = {};
+  materializeColumnAsColumnarArray(schema, columnData, rowGroup.rowCount, key, columns);
+  return columns[schema.findFieldBranch(key)[0].name];
+}
+
 // eslint-disable-next-line max-statements, complexity
 function materializeColumnAsColumnarArray(
   schema: ParquetSchema,
@@ -300,23 +357,12 @@ function materializeColumnAsColumnarArray(
 
   const columnName = branch[0].name;
 
-  let column: ArrayType | undefined;
-  const {values} = columnData;
-  if (values.length === rowCount && branch[0].primitiveType) {
-    // if (branch[0].repetitionType === `REQUIRED`) {
-    //   switch (branch[0].primitiveType) {
-    //     case 'INT32': return values instanceof Int32Array ? values : new Int32Array(values);
-    //   }
-    // }
-    column = values;
-  }
-
-  if (column) {
-    columns[columnName] = column;
+  if (branch.length === 1 && field.repetitionType !== 'REPEATED') {
+    columns[columnName] = materializeFlatColumn(field, columnData, rowCount);
     return;
   }
 
-  column = new Array(rowCount);
+  const column: ArrayType = new Array(rowCount);
   for (let i = 0; i < rowCount; i++) {
     column[i] = {};
   }
@@ -397,3 +443,34 @@ function materializeColumnAsColumnarArray(
     }
   }
 }
+
+/** Materializes a required or optional top-level primitive column with logical type conversion. */
+function materializeFlatColumn(
+  field: ParquetField,
+  columnData: ParquetColumnChunk,
+  rowCount: number
+): ArrayType {
+  const logicalType = field.originalType || field.primitiveType!;
+  const fromPrimitive = Types.PARQUET_LOGICAL_TYPES[logicalType].fromPrimitive;
+  const count = Math.min(columnData.count, rowCount);
+
+  if (
+    field.repetitionType === 'REQUIRED' &&
+    !fromPrimitive &&
+    count === rowCount &&
+    columnData.values.length === rowCount
+  ) {
+    return columnData.values;
+  }
+
+  const column = new Array(rowCount).fill(null);
+  let valueIndex = 0;
+  for (let rowIndex = 0; rowIndex < count; rowIndex++) {
+    if (columnData.dlevels[rowIndex] === field.dLevelMax) {
+      const primitiveValue = columnData.values[valueIndex];
+      column[rowIndex] = fromPrimitive ? fromPrimitive(primitiveValue, field) : primitiveValue;
+      valueIndex++;
+    }
+  }
+  return column;
+}