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

package/src/lib/parquet-page-index.ts

@@ -32,7 +32,11 @@ import {
 } from '../parquetjs/utils/binary-utils';
 import {Uint8ArrayCompactProtocol} from '../parquetjs/utils/uint8-array-compact-protocol';
 import {Uint8ArrayTransport} from '../parquetjs/utils/uint8-array-transport';
-import {canParquetStatisticsMatch, getParquetPredicateColumns} from './parquet-predicate';
+import {
+  canParquetStatisticsMatch,
+  getParquetPredicatePath,
+  getParquetPredicatePaths
+} from './parquet-predicate';
 
 /** Half-open logical row range relative to one Parquet row group. */
 export type ParquetRowRange = {
@@ -73,7 +77,27 @@ export type ParquetPagePruningPlan = {
   prunedRowCount: number;
 };
 
-type ParquetPageStatistics = {
+/** Decrypts an encrypted page-index module before Thrift decoding. */
+export type ParquetPageIndexDecryptor = (
+  bytes: Uint8Array,
+  module: 'column-index' | 'offset-index',
+  rowGroupOrdinal: number,
+  columnOrdinal: number,
+  columnChunk: ColumnChunk
+) => Promise<Uint8Array>;
+
+/** Optional controls for reading page indexes from a Parquet source. */
+export type ParquetPagePruningOptions = {
+  /** Row-group ordinal used when constructing encrypted-module AAD. */
+  rowGroupOrdinal?: number;
+  /** Decode legacy INT96 page statistics as epoch nanoseconds. */
+  int96AsTimestamp?: boolean;
+  /** Decrypts encrypted index modules; omitted for plaintext files. */
+  decryptModule?: ParquetPageIndexDecryptor;
+};
+
+/** Logical statistics attached to one data page. */
+export type ParquetPageStatistics = {
   min?: unknown;
   max?: unknown;
   nullCount?: number;
@@ -89,8 +113,10 @@ type ParquetColumnPageStatistics = {
 /**
  * Builds a conservative selective-page plan from Parquet column and offset indexes.
  *
- * Returns `undefined` when indexes or the current flat-column decoder cannot safely avoid full
- * column-chunk reads. An empty `rowRanges` array means the indexes prove the predicate cannot match.
+ * Returns `undefined` when indexes or the selected decoder cannot safely avoid full column-chunk
+ * reads. Repeated leaves are selected only when every decoded column uses the same page row
+ * boundaries, so ranges begin and end on complete logical rows. An empty `rowRanges` array means
+ * the indexes prove the predicate cannot match.
  */
 export async function createParquetPagePruningPlan(
   file: ReadableFile,
@@ -98,23 +124,27 @@ export async function createParquetPagePruningPlan(
   schema: ParquetSchema,
   selectedColumnPaths: readonly string[][],
   predicate: ParquetPredicate,
-  signal?: AbortSignal
+  signal?: AbortSignal,
+  options?: ParquetPagePruningOptions
 ): Promise<ParquetPagePruningPlan | undefined> {
   const rowCount = Number(rowGroup.num_rows);
   const selectedColumnChunks = getSelectedColumnChunks(rowGroup, selectedColumnPaths);
-  if (!isSafeFlatPageSelection(schema, selectedColumnChunks) || rowCount <= 0) {
+  if (!isSafePageSelection(schema, selectedColumnChunks) || rowCount <= 0) {
     return undefined;
   }
 
   const pageLocations: ParquetPageLocations = {};
   const pageStatistics: Record<string, ParquetColumnPageStatistics> = {};
-  const predicateColumns = new Set(getParquetPredicateColumns(predicate));
+  const predicatePaths = new Set(
+    getParquetPredicatePaths(predicate).map(path => JSON.stringify(path))
+  );
   let indexCount = 0;
 
   await Promise.all(
     selectedColumnChunks.map(async columnChunk => {
+      const columnOrdinal = rowGroup.columns.indexOf(columnChunk);
       const path = columnChunk.meta_data!.path_in_schema;
-      const pathKey = path.join('.');
+      const pathKey = JSON.stringify(path);
       const offsetIndexRange = getParquetIndexRange(
         columnChunk.offset_index_offset,
         columnChunk.offset_index_length,
@@ -130,14 +160,23 @@ export async function createParquetPagePruningPlan(
       );
       let pages: ParquetDataPageLocation[];
       try {
-        pages = decodeOffsetIndex(toUint8Array(offsetIndexBytes), rowCount);
+        const decryptedBytes = options?.decryptModule
+          ? await options.decryptModule(
+              toUint8Array(offsetIndexBytes),
+              'offset-index',
+              options.rowGroupOrdinal ?? 0,
+              columnOrdinal,
+              columnChunk
+            )
+          : toUint8Array(offsetIndexBytes);
+        pages = decodeOffsetIndex(decryptedBytes, rowCount);
       } catch {
         return;
       }
       pageLocations[pathKey] = pages;
       indexCount++;
 
-      if (path.length !== 1 || !predicateColumns.has(path[0])) {
+      if (!predicatePaths.has(pathKey)) {
         return;
       }
       const columnIndexRange = getParquetIndexRange(
@@ -155,9 +194,18 @@ export async function createParquetPagePruningPlan(
       );
       const field = schema.findField(path);
       try {
-        pageStatistics[path[0]] = {
+        const decryptedBytes = options?.decryptModule
+          ? await options.decryptModule(
+              toUint8Array(columnIndexBytes),
+              'column-index',
+              options.rowGroupOrdinal ?? 0,
+              columnOrdinal,
+              columnChunk
+            )
+          : toUint8Array(columnIndexBytes);
+        pageStatistics[pathKey] = {
           pages,
-          statistics: decodeColumnIndex(toUint8Array(columnIndexBytes), pages, field)
+          statistics: decodeParquetColumnIndex(decryptedBytes, pages, field, options)
         };
       } catch {
         return;
@@ -168,12 +216,16 @@ export async function createParquetPagePruningPlan(
 
   if (
     selectedColumnChunks.some(
-      columnChunk => !pageLocations[columnChunk.meta_data!.path_in_schema.join('.')]
+      columnChunk => !pageLocations[JSON.stringify(columnChunk.meta_data!.path_in_schema)]
     )
   ) {
     return undefined;
   }
 
+  if (!hasCompatiblePageBoundaries(schema, selectedColumnChunks, pageLocations)) {
+    return undefined;
+  }
+
   const candidateRanges = getPredicateRowRanges(predicate, pageStatistics, rowCount);
   if (candidateRanges === undefined) {
     return undefined;
@@ -206,18 +258,21 @@ export async function createParquetPagePruningPlan(
 export function getParquetPageReadRanges(
   rowGroup: RowGroup,
   selectedColumnPaths: readonly string[][],
-  plan: ParquetPagePruningPlan
+  plan: ParquetPagePruningPlan,
+  schema: ParquetSchema
 ): Array<{offset: number; length: number}> {
   const ranges: Array<{offset: number; length: number}> = [];
   for (const columnChunk of getSelectedColumnChunks(rowGroup, selectedColumnPaths)) {
     const columnMetadata = columnChunk.meta_data!;
-    const pages = plan.pageLocations[columnMetadata.path_in_schema.join('.')];
+    const pages = plan.pageLocations[JSON.stringify(columnMetadata.path_in_schema)];
     const selectedPages = pages.filter(page =>
       plan.rowRanges.some(range => page.endRowIndex > range.start && page.firstRowIndex < range.end)
     );
     if (!selectedPages.length) {
       continue;
     }
+    const field = schema.findField(columnMetadata.path_in_schema);
+    const repeated = field.rLevelMax > 0 || field.repetitionType === 'REPEATED';
     const dictionaryPageOffset = Number(columnMetadata.dictionary_page_offset);
     if (Number.isSafeInteger(dictionaryPageOffset) && dictionaryPageOffset > 0) {
       ranges.push({
@@ -232,7 +287,9 @@ export function getParquetPageReadRanges(
       if (!overlappingPages.length) {
         continue;
       }
-      const firstPage = overlappingPages[0];
+      // A repeated page may begin with a continuation level. Transfer the complete prefix so
+      // worker-backed readers can probe predecessor pages before decoding the selected range.
+      const firstPage = repeated ? pages[0] : overlappingPages[0];
       const lastPage = overlappingPages[overlappingPages.length - 1];
       ranges.push({
         offset: firstPage.offset,
@@ -253,10 +310,27 @@ export function decodeOffsetIndex(bytes: Uint8Array, rowCount: number): ParquetD
   ) {
     throw new Error('Invalid Parquet offset index page locations');
   }
+  let previousFirstRowIndex = -1;
+  for (const location of offsetIndex.page_locations) {
+    const firstRowIndex = Number(location.first_row_index);
+    if (!Number.isSafeInteger(firstRowIndex) || firstRowIndex < previousFirstRowIndex) {
+      throw new Error('Invalid Parquet offset index page locations');
+    }
+    previousFirstRowIndex = firstRowIndex;
+  }
   return offsetIndex.page_locations.map((location, index, locations) => {
     const firstRowIndex = Number(location.first_row_index);
-    const endRowIndex =
-      index + 1 < locations.length ? Number(locations[index + 1].first_row_index) : rowCount;
+    // A repeated row may continue on one or more pages. In that case the offset index
+    // legitimately repeats first_row_index. Extend each continuation page to the next
+    // strictly larger row start so selecting the row includes the complete continuation.
+    let endRowIndex = rowCount;
+    for (let nextIndex = index + 1; nextIndex < locations.length; nextIndex++) {
+      const nextFirstRowIndex = Number(locations[nextIndex].first_row_index);
+      if (nextFirstRowIndex > firstRowIndex) {
+        endRowIndex = nextFirstRowIndex;
+        break;
+      }
+    }
     const offset = Number(location.offset);
     const compressedByteLength = location.compressed_page_size;
     if (
@@ -308,7 +382,7 @@ function getPredicateRowRanges(
     | ParquetComparisonPredicate
     | ParquetInPredicate
     | ParquetNullPredicate;
-  const column = columns[leafPredicate.args[0].property];
+  const column = columns[JSON.stringify(getParquetPredicatePath(leafPredicate.args[0]))];
   if (!column) {
     return undefined;
   }
@@ -328,11 +402,12 @@ function getPredicateRowRanges(
   return mergeRowRanges(ranges);
 }
 
-/** Decodes column-index values into the logical representation used by exact predicates. */
-function decodeColumnIndex(
+/** Decodes column-index values into logical per-page statistics. */
+export function decodeParquetColumnIndex(
   bytes: Uint8Array,
   pages: readonly ParquetDataPageLocation[],
-  field: ParquetField
+  field: ParquetField,
+  options?: Pick<ParquetPagePruningOptions, 'int96AsTimestamp'>
 ): ParquetPageStatistics[] {
   const protocol = new Uint8ArrayCompactProtocol(new Uint8ArrayTransport(bytes));
   const columnIndex = ColumnIndex.read(protocol as any);
@@ -347,10 +422,10 @@ function decodeColumnIndex(
   return pages.map((_page, pageIndex) => ({
     min: columnIndex.null_pages[pageIndex]
       ? undefined
-      : decodeParquetPageStatisticsValue(columnIndex.min_values[pageIndex], field),
+      : decodeParquetPageStatisticsValue(columnIndex.min_values[pageIndex], field, options),
     max: columnIndex.null_pages[pageIndex]
       ? undefined
-      : decodeParquetPageStatisticsValue(columnIndex.max_values[pageIndex], field),
+      : decodeParquetPageStatisticsValue(columnIndex.max_values[pageIndex], field, options),
     nullCount:
       columnIndex.null_counts?.[pageIndex] !== undefined
         ? Number(columnIndex.null_counts[pageIndex])
@@ -363,7 +438,11 @@ function decodeColumnIndex(
 }
 
 /** Decodes one physical page-index statistic and applies its Parquet logical annotation. */
-export function decodeParquetPageStatisticsValue(bytes: Uint8Array, field: ParquetField): unknown {
+export function decodeParquetPageStatisticsValue(
+  bytes: Uint8Array,
+  field: ParquetField,
+  options?: Pick<ParquetPagePruningOptions, 'int96AsTimestamp'>
+): unknown {
   let primitiveValue: unknown;
   switch (field.primitiveType) {
     case 'BOOLEAN':
@@ -384,6 +463,18 @@ export function decodeParquetPageStatisticsValue(bytes: Uint8Array, field: Parqu
       primitiveValue = readDoubleLE(bytes, 0);
       break;
     case 'INT96':
+      if (options?.int96AsTimestamp && bytes.byteLength >= 12) {
+        const dataView = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+        const nanosecondsOfDay = dataView.getBigUint64(0, true);
+        if (nanosecondsOfDay < 86_400_000_000_000n) {
+          primitiveValue =
+            (BigInt(dataView.getInt32(8, true)) - 2440588n) * 86_400_000_000_000n +
+            nanosecondsOfDay;
+          break;
+        }
+      }
+      primitiveValue = copyUint8Array(bytes);
+      break;
     case 'BYTE_ARRAY':
     case 'FIXED_LEN_BYTE_ARRAY':
       primitiveValue = copyUint8Array(bytes);
@@ -409,16 +500,13 @@ function getSelectedColumnChunks(
   });
 }
 
-/** Restricts selective page reads to independently materializable flat primitive columns. */
-function isSafeFlatPageSelection(
-  schema: ParquetSchema,
-  columnChunks: readonly ColumnChunk[]
-): boolean {
+/** Restricts selective page reads to primitive leaves with safe page boundaries. */
+function isSafePageSelection(schema: ParquetSchema, columnChunks: readonly ColumnChunk[]): boolean {
   return (
     columnChunks.length > 0 &&
     columnChunks.every(columnChunk => {
       const path = columnChunk.meta_data?.path_in_schema;
-      if (!path || path.length !== 1) {
+      if (!path) {
         return false;
       }
       const field = schema.findField(path);
@@ -427,9 +515,7 @@ function isSafeFlatPageSelection(
       );
       const dictionaryPageOffset = Number(columnChunk.meta_data!.dictionary_page_offset);
       return (
-        field.path.length === 1 &&
-        field.repetitionType !== 'REPEATED' &&
-        field.rLevelMax === 0 &&
+        field.primitiveType !== undefined &&
         (!dictionaryEncoded ||
           (Number.isSafeInteger(dictionaryPageOffset) && dictionaryPageOffset > 0))
       );
@@ -437,6 +523,66 @@ function isSafeFlatPageSelection(
   );
 }
 
+/**
+ * Returns whether a column can be decoded from selected page ranges without changing row shape.
+ *
+ * Repeated leaves are safe only when the selected columns share complete logical-row boundaries.
+ */
+export function canUseParquetPageIndexForColumn(
+  schema: ParquetSchema,
+  columnChunk: ColumnChunk
+): boolean {
+  const path = columnChunk.meta_data?.path_in_schema;
+  if (!path) {
+    return false;
+  }
+  const field = schema.findField(path);
+  const dictionaryEncoded = columnChunk.meta_data!.encodings.some(
+    encoding => encoding === Encoding.PLAIN_DICTIONARY || encoding === Encoding.RLE_DICTIONARY
+  );
+  const dictionaryPageOffset = Number(columnChunk.meta_data!.dictionary_page_offset);
+  return (
+    field.primitiveType !== undefined &&
+    field.repetitionType !== 'REPEATED' &&
+    field.rLevelMax === 0 &&
+    (!dictionaryEncoded || (Number.isSafeInteger(dictionaryPageOffset) && dictionaryPageOffset > 0))
+  );
+}
+
+/** Ensures repeated selective reads keep every selected column on identical row boundaries. */
+function hasCompatiblePageBoundaries(
+  schema: ParquetSchema,
+  columnChunks: readonly ColumnChunk[],
+  pageLocations: ParquetPageLocations
+): boolean {
+  const hasRepeatedLeaf = columnChunks.some(columnChunk => {
+    const field = schema.findField(columnChunk.meta_data!.path_in_schema);
+    return field.rLevelMax > 0 || field.repetitionType === 'REPEATED';
+  });
+  if (!hasRepeatedLeaf) {
+    return true;
+  }
+  // Offset indexes identify the first logical row represented by each page, but do not
+  // expose repetition levels. Equal starts therefore indicate a continuation page whose
+  // boundary cannot be proven safe for a selective read; retain the full-column fallback.
+  if (
+    columnChunks.some(columnChunk => {
+      const pages = pageLocations[JSON.stringify(columnChunk.meta_data!.path_in_schema)];
+      return pages.some(
+        (page, pageIndex) =>
+          pageIndex > 0 && page.firstRowIndex === pages[pageIndex - 1].firstRowIndex
+      );
+    })
+  ) {
+    return false;
+  }
+  const signatures = columnChunks.map(columnChunk => {
+    const pages = pageLocations[JSON.stringify(columnChunk.meta_data!.path_in_schema)];
+    return pages.map(page => `${page.firstRowIndex}:${page.endRowIndex}`).join('|');
+  });
+  return signatures.every(signature => signature === signatures[0]);
+}
+
 /** Validates one optional footer index byte range against the containing file. */
 export function getParquetIndexRange(
   offsetValue: ColumnChunk['offset_index_offset'],

package/src/lib/parquet-predicate.ts

@@ -12,6 +12,7 @@ import type {
   ParquetNotPredicate,
   ParquetNullPredicate,
   ParquetPredicate,
+  ParquetPredicateProperty,
   ParquetPredicateValue,
   ParquetRowGroupMetadata
 } from '../parquet-source-types';
@@ -19,14 +20,20 @@ import {encodeUtf8} from '../parquetjs/utils/binary-utils';
 
 /** Returns the unique top-level columns referenced by a predicate. */
 export function getParquetPredicateColumns(predicate: ParquetPredicate): string[] {
-  const columns = new Set<string>();
+  return [...new Set(getParquetPredicatePaths(predicate).map(path => path[0]))];
+}
+
+/** Returns the unique physical schema paths referenced by a predicate. */
+export function getParquetPredicatePaths(predicate: ParquetPredicate): string[][] {
+  const paths = new Map<string, string[]>();
   visitParquetPredicate(predicate, child => {
     const property = getPredicateProperty(child);
     if (property) {
-      columns.add(property.property);
+      const path = [...getParquetPredicatePath(property)];
+      paths.set(path.join('\0'), path);
     }
   });
-  return [...columns];
+  return [...paths.values()];
 }
 
 /** Copies a predicate tree and mutable scalar values for read-scoped option snapshots. */
@@ -75,8 +82,12 @@ export function validateParquetPredicate(
       return;
     }
     const property = child.args[0];
-    if (!availableColumns.has(property.property)) {
-      throw new Error(`Parquet predicate column not found: ${property.property}`);
+    const path = getParquetPredicatePath(property);
+    if (!path.length || path.some(pathComponent => !pathComponent)) {
+      throw new Error('Parquet predicate property path must contain non-empty strings');
+    }
+    if (!availableColumns.has(path[0])) {
+      throw new Error(`Parquet predicate column not found: ${path.join('.')}`);
     }
     if (child.op === 'in' && child.args[1].length === 0) {
       throw new Error('Parquet predicate in requires at least one value');
@@ -99,11 +110,11 @@ export function canParquetRowGroupMatch(
     if (child.op !== 'isNull') {
       return true;
     }
-    const column = getRowGroupColumn(rowGroup, child.args[0].property);
+    const column = getRowGroupColumn(rowGroup, getParquetPredicatePath(child.args[0]));
     return !column?.statistics || column.statistics.nullCount !== rowGroup.rowCount;
   }
 
-  const column = getRowGroupColumn(rowGroup, predicate.args[0].property);
+  const column = getRowGroupColumn(rowGroup, getParquetPredicatePath(predicate.args[0]));
   const statistics = column?.statistics;
   if (!statistics) {
     return true;
@@ -160,7 +171,7 @@ export function gatherParquetColumns(
     if (selectedColumns && !selectedColumns.has(name)) {
       continue;
     }
-    gatheredColumns[name] = rowIndices.map(rowIndex => getColumnValue(column, rowIndex));
+    gatheredColumns[name] = rowIndices.map(rowIndex => (column as ArrayLike<unknown>)[rowIndex]);
   }
   return gatheredColumns;
 }
@@ -182,7 +193,7 @@ function evaluateParquetPredicate(
     const result = evaluateParquetPredicate(predicate.args[0], columns, rowIndex);
     return result === null ? null : !result;
   }
-  const value = getColumnValue(columns[predicate.args[0].property], rowIndex);
+  const value = getColumnValue(columns, getParquetPredicatePath(predicate.args[0]), rowIndex);
   if (predicate.op === 'isNull') {
     return value === null || value === undefined;
   }
@@ -311,11 +322,27 @@ function compareUint8Arrays(left: Uint8Array, right: Uint8Array): number {
 }
 
 /** Reads one value from an array-like materialized Parquet column. */
-function getColumnValue(column: ArrayType | undefined, rowIndex: number): unknown {
+function getColumnValue(
+  columns: Record<string, ArrayType>,
+  path: readonly string[],
+  rowIndex: number
+): unknown {
+  const column = columns[path[0]];
   if (!column) {
     throw new Error('Parquet predicate column was not decoded');
   }
-  return (column as ArrayLike<unknown>)[rowIndex];
+  let value = (column as ArrayLike<unknown>)[rowIndex];
+  for (let pathIndex = 1; pathIndex < path.length; pathIndex++) {
+    if (value === null || value === undefined || typeof value !== 'object') {
+      return undefined;
+    }
+    const property = path[pathIndex];
+    value =
+      typeof (value as {get?: unknown}).get === 'function'
+        ? (value as {get: (name: string) => unknown}).get(property)
+        : (value as Record<string, unknown>)[property];
+  }
+  return value;
 }
 
 /** Copies mutable binary and date predicate values. */
@@ -330,21 +357,30 @@ function copyPredicateValue(value: ParquetPredicateValue): ParquetPredicateValue
 }
 
 /** Copies a predicate property reference. */
-function copyPredicateProperty(property: {property: string}): {property: string} {
-  return {property: property.property};
+function copyPredicateProperty(property: ParquetPredicateProperty): ParquetPredicateProperty {
+  return {
+    property: Array.isArray(property.property) ? [...property.property] : property.property
+  };
+}
+
+/** Returns one predicate property as an explicit Parquet schema path. */
+export function getParquetPredicatePath(property: ParquetPredicateProperty): readonly string[] {
+  return Array.isArray(property.property) ? property.property : [property.property as string];
 }
 
 /** Returns the property reference from a leaf predicate. */
-function getPredicateProperty(predicate: ParquetPredicate): {property: string} | undefined {
+function getPredicateProperty(predicate: ParquetPredicate): ParquetPredicateProperty | undefined {
   return isParquetLogicalPredicate(predicate) || isParquetNotPredicate(predicate)
     ? undefined
     : predicate.args[0];
 }
 
 /** Returns one top-level column chunk from normalized row-group metadata. */
-function getRowGroupColumn(rowGroup: ParquetRowGroupMetadata, property: string) {
+function getRowGroupColumn(rowGroup: ParquetRowGroupMetadata, path: readonly string[]) {
   return rowGroup.columns.find(
-    columnChunk => columnChunk.path.length === 1 && columnChunk.path[0] === property
+    columnChunk =>
+      columnChunk.path.length === path.length &&
+      columnChunk.path.every((pathComponent, index) => pathComponent === path[index])
   );
 }