@loaders.gl/parquet 4.5.3 → 5.0.0-alpha.10

package/src/parquetjs/schema/shred.ts

@@ -8,6 +8,7 @@ import {ArrayType} from '@loaders.gl/schema';
 import {ParquetRowGroup, ParquetColumnChunk, ParquetField, ParquetRow} from './declare';
 import {ParquetSchema} from './schema';
 import * as Types from './types';
+import {decodeVariant} from './variant';
 
 export {ParquetRowGroup};
 
@@ -65,9 +66,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;
   }
 }
@@ -91,10 +101,33 @@ function shredRecordFields(
       record[field.name] !== undefined &&
       record[field.name] !== null
     ) {
-      if (record[field.name].constructor === Array) {
-        values = record[field.name];
+      const fieldValue = record[field.name];
+      if (field.logicalType?.type === 'LIST') {
+        // Normalize the high-level Arrow/list value to Parquet's standard
+        // three-level LIST representation before descending into the wrapper.
+        const listValues = Array.isArray(fieldValue) ? fieldValue : [];
+        if (field.fields?.list) {
+          values.push({list: listValues.map(element => ({element}))});
+        } else {
+          // Accept the legacy two-level LIST layout where the repeated field
+          // itself owns the element leaf.
+          values = listValues;
+        }
+      } else if (field.logicalType?.type === 'MAP') {
+        // Normalize Map/object/entry-array values to the standard MAP_KEY_VALUE
+        // wrapper expected by the shredding algorithm.
+        const mapEntries = normalizeMapEntries(fieldValue);
+        if (field.fields?.key_value) {
+          values.push({key_value: mapEntries});
+        } else {
+          // Accept legacy map layouts whose key/value group is the repeated
+          // field itself.
+          values = mapEntries;
+        }
+      } else if (fieldValue.constructor === Array) {
+        values = fieldValue;
       } else {
-        values.push(record[field.name]);
+        values.push(fieldValue);
       }
     }
     // check values
@@ -111,8 +144,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,16 +157,38 @@ 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)
         );
       }
     }
   }
 }
 
+/** Converts supported JavaScript map representations to Parquet map entries. */
+function normalizeMapEntries(value: unknown): Array<{key: unknown; value: unknown}> {
+  if (value instanceof Map) {
+    return Array.from(value, ([key, mapValue]) => ({key, value: mapValue}));
+  }
+  if (Array.isArray(value)) {
+    return value.flatMap(entry => {
+      if (Array.isArray(entry) && entry.length >= 2) {
+        return [{key: entry[0], value: entry[1]}];
+      }
+      if (entry && typeof entry === 'object') {
+        return [{key: Reflect.get(entry, 'key'), value: Reflect.get(entry, 'value')}];
+      }
+      return [];
+    });
+  }
+  if (value && typeof value === 'object') {
+    return Object.entries(value).map(([key, mapValue]) => ({key, value: mapValue}));
+  }
+  return [];
+}
+
 /**
  * 'Materialize' a list of <value, repetition_level, definition_level>
  * tuples back to nested records (objects/arrays) using the Google Dremel
@@ -154,10 +209,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];
@@ -165,9 +219,111 @@ export function materializeRows(schema: ParquetSchema, rowGroup: ParquetRowGroup
       materializeColumnAsRows(schema, columnData, key, rows);
     }
   }
+  decodeVariantFields(schema.fields, rows);
   return rows;
 }
 
+/** Decodes complete unshredded VARIANT groups after their binary child columns are materialized. */
+function decodeVariantFields(fields: Record<string, ParquetField>, records: unknown[]): void {
+  for (const field of Object.values(fields)) {
+    if (field.logicalType?.type === 'VARIANT') {
+      for (const record of records) {
+        if (record && typeof record === 'object') {
+          const row = record as Record<string, unknown>;
+          if (Object.prototype.hasOwnProperty.call(row, field.name)) {
+            row[field.name] = decodeVariantRecord(row[field.name], field);
+          }
+        }
+      }
+    } else if (field.fields) {
+      for (const record of records) {
+        if (record && typeof record === 'object') {
+          decodeVariantFields(field.fields, getNestedRecords(record, field.name));
+        }
+      }
+    }
+  }
+}
+
+/** Recursively decodes one materialized Variant group, preserving shredded values as records. */
+function decodeVariantRecord(value: unknown, field?: ParquetField): unknown {
+  if (Array.isArray(value)) {
+    return value.map(item => decodeVariantRecord(item, field));
+  }
+  if (!value || typeof value !== 'object') {
+    return value;
+  }
+  const record = value as Record<string, unknown>;
+  const metadata = record.metadata;
+  const variantValue = record.value;
+  if (isByteArray(metadata) && isByteArray(variantValue)) {
+    return decodeVariant(metadata, variantValue);
+  }
+  // Shredded Variant encodings may materialize the typed_value child directly while
+  // omitting the unshredded metadata/value pair. Preserve that typed representation as
+  // the logical value instead of leaking the physical wrapper into object rows.
+  if (Object.prototype.hasOwnProperty.call(record, 'typed_value')) {
+    return decodeShreddedVariantValue(record.typed_value, field?.fields?.typed_value);
+  }
+  return value;
+}
+
+/** Collapses a declared one-field typed-value union used by shredded Variant layouts. */
+function decodeShreddedVariantValue(value: unknown, field?: ParquetField): unknown {
+  if (!value || typeof value !== 'object' || Array.isArray(value)) {
+    return value;
+  }
+  const record = value as Record<string, unknown>;
+  const keys = Object.keys(record).filter(key => record[key] !== undefined);
+  if (keys.length === 1 && isVariantTypedValueField(field?.fields?.[keys[0]])) {
+    return decodeShreddedVariantValue(record[keys[0]], field?.fields?.[keys[0]]);
+  }
+  return value;
+}
+
+/** Identifies the declared members of the Parquet Variant typed-value union. */
+function isVariantTypedValueField(field: ParquetField | undefined): boolean {
+  return Boolean(field && VARIANT_TYPED_VALUE_FIELD_NAMES.has(field.name));
+}
+
+const VARIANT_TYPED_VALUE_FIELD_NAMES = new Set([
+  'null_value',
+  'boolean_value',
+  'int8_value',
+  'int16_value',
+  'int32_value',
+  'int64_value',
+  'uint8_value',
+  'uint16_value',
+  'uint32_value',
+  'uint64_value',
+  'float16_value',
+  'float32_value',
+  'float64_value',
+  'decimal4_value',
+  'decimal8_value',
+  'decimal16_value',
+  'string_value',
+  'blob_value',
+  'date_value',
+  'timestamp_millis_value',
+  'timestamp_micros_value',
+  'timestamp_nanos_value',
+  'object_value',
+  'array_value'
+]);
+
+/** Finds nested records while preserving repeated group values. */
+function getNestedRecords(record: object, fieldName: string): unknown[] {
+  const value = Reflect.get(record, fieldName);
+  return Array.isArray(value) ? value : value && typeof value === 'object' ? [value] : [];
+}
+
+/** Identifies binary values emitted by Parquet primitive materialization. */
+function isByteArray(value: unknown): value is ArrayBuffer | ArrayBufferView {
+  return value instanceof ArrayBuffer || ArrayBuffer.isView(value);
+}
+
 /** Populate record fields for one column */
 // eslint-disable-next-line max-statements, complexity
 function materializeColumnAsRows(
@@ -179,17 +335,32 @@ 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;
   for (let i = 0; i < columnData.count; i++) {
-    const dLevel = columnData.dlevels[i];
-    const rLevel = columnData.rlevels[i];
+    // Zero maximum levels are implicit and may use allocation-free empty buffers. Preserve both
+    // fallbacks: required nested scalars still need dLevel === 0 to materialize their leaf value.
+    const dLevel = columnData.dlevels[i] ?? 0;
+    const rLevel = columnData.rlevels[i] ?? 0;
     rLevels[rLevel]++;
     rLevels.fill(0, rLevel + 1);
 
     let rIndex = 0;
     let record = rows[rLevels[rIndex++] - 1];
+    if (!record) {
+      throw new Error(
+        `Parquet column ${key} referenced row ${rLevels[0] - 1} of ${rows.length} at value ${i} (rLevel ${rLevel}, dLevel ${dLevel})`
+      );
+    }
 
     // Internal nodes - Build a nested row object
     for (const step of branch) {
@@ -219,12 +390,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 = getDecodedColumnValue(columnData, vIndex);
+      const value = fromPrimitive ? fromPrimitive(primitiveValue, field) : primitiveValue;
       vIndex++;
 
       switch (field.repetitionType) {
@@ -248,6 +415,45 @@ 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] = getDecodedColumnValue(columnData, rowIndex);
+    }
+    return;
+  }
+
+  if (field.repetitionType === 'REQUIRED') {
+    for (let rowIndex = 0; rowIndex < count; rowIndex++) {
+      rows[rowIndex][field.name] = fromPrimitive!(
+        getDecodedColumnValue(columnData, rowIndex),
+        field
+      );
+    }
+    return;
+  }
+
+  for (let rowIndex = 0; rowIndex < count; rowIndex++) {
+    if (columnData.dlevels[rowIndex] === field.dLevelMax) {
+      const primitiveValue = getDecodedColumnValue(columnData, valueIndex);
+      rows[rowIndex][field.name] = fromPrimitive
+        ? fromPrimitive(primitiveValue, field)
+        : primitiveValue;
+      valueIndex++;
+    }
+  }
+}
+
 // Columnar export
 
 /**
@@ -275,14 +481,72 @@ 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) {
+      const columnName = schema.findFieldBranch(key)[0].name;
+      columns[columnName] = columns[columnName]
+        ? mergeMaterializedColumn(columns[columnName], column)
+        : column;
     }
   }
   return columns;
 }
 
+/** Merges independently decoded nested leaf columns into one top-level column. */
+function mergeMaterializedColumn(existing: ArrayType, incoming: ArrayType): ArrayType {
+  const length = Math.max(existing.length, incoming.length);
+  const merged = new Array(length);
+  for (let index = 0; index < length; index++) {
+    merged[index] = mergeMaterializedValue(existing[index], incoming[index]);
+  }
+  return merged;
+}
+
+/** Recursively combines objects and corresponding repeated elements without losing scalar leaves. */
+function mergeMaterializedValue(existing: unknown, incoming: unknown): unknown {
+  if (existing === null || existing === undefined) {
+    return incoming;
+  }
+  if (incoming === null || incoming === undefined) {
+    return existing;
+  }
+  if (Array.isArray(existing) && Array.isArray(incoming)) {
+    const length = Math.max(existing.length, incoming.length);
+    return Array.from({length}, (_, index) =>
+      mergeMaterializedValue(existing[index], incoming[index])
+    );
+  }
+  if (
+    typeof existing === 'object' &&
+    typeof incoming === 'object' &&
+    !ArrayBuffer.isView(existing) &&
+    !ArrayBuffer.isView(incoming)
+  ) {
+    const merged: Record<string, unknown> = {...(existing as Record<string, unknown>)};
+    for (const [key, value] of Object.entries(incoming as Record<string, unknown>)) {
+      merged[key] = mergeMaterializedValue(merged[key], value);
+    }
+    return merged;
+  }
+  return incoming;
+}
+
+/** 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 +564,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] = {};
   }
@@ -365,7 +618,7 @@ function materializeColumnAsColumnarArray(
       const value = Types.fromPrimitive(
         // @ts-ignore
         field.originalType || field.primitiveType,
-        columnData.values[vIndex],
+        getDecodedColumnValue(columnData, vIndex),
         field
       );
       vIndex++;
@@ -397,3 +650,54 @@ 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;
+  }
+
+  if (field.repetitionType === 'REQUIRED') {
+    const column = new Array(rowCount).fill(null);
+    for (let rowIndex = 0; rowIndex < count; rowIndex++) {
+      column[rowIndex] = fromPrimitive
+        ? fromPrimitive(getDecodedColumnValue(columnData, rowIndex), field)
+        : getDecodedColumnValue(columnData, rowIndex);
+    }
+    return column;
+  }
+
+  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 = getDecodedColumnValue(columnData, valueIndex);
+      column[rowIndex] = fromPrimitive ? fromPrimitive(primitiveValue, field) : primitiveValue;
+      valueIndex++;
+    }
+  }
+  return column;
+}
+
+/** Returns one compact byte value or one value from the decoder's conventional destination. */
+function getDecodedColumnValue(columnData: ParquetColumnChunk, valueIndex: number): unknown {
+  const byteArrayData = columnData.byteArrayData;
+  if (!byteArrayData) return columnData.values[valueIndex];
+  return byteArrayData.data.subarray(
+    byteArrayData.valueOffsets[valueIndex],
+    byteArrayData.valueOffsets[valueIndex + 1]
+  );
+}