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

package/src/parquetjs/compression.ts

@@ -5,96 +5,66 @@
 // Forked from https://github.com/kbajalc/parquets under MIT license
 // Forked from https://github.com/ironSource/parquetjs under MIT license
 
+import type {Compression} from '@loaders.gl/compression';
 import {
-  Compression,
-  NoCompression,
-  GZipCompression,
-  SnappyCompression,
-  BrotliCompression,
-  // LZOCompression,
-  LZ4Compression,
-  ZstdCompression
-} from '@loaders.gl/compression';
-import {registerJSModules} from '@loaders.gl/loader-utils';
+  decompressWithNativeDecompressionStream,
+  type NativeDecompressionFormat
+} from '@loaders.gl/compression/native-decompression';
+import {getJSModuleOrNull, registerJSModules} from '@loaders.gl/loader-utils';
 
 import {ParquetCompression} from './schema/declare';
-
-/** We can't use loaders-util buffer handling since we are dependent on buffers even in the browser */
-function toBuffer(arrayBuffer: ArrayBuffer): Buffer {
-  return Buffer.from(arrayBuffer);
-}
-
-function toArrayBuffer(buffer: Buffer): ArrayBuffer {
-  // TODO - per docs we should just be able to call buffer.buffer, but there are issues
-  if (Buffer.isBuffer(buffer)) {
-    const typedArray = new Uint8Array(buffer.buffer, buffer.byteOffset, buffer.length);
-    return typedArray.slice().buffer;
-  }
-  return buffer;
-}
-
-// TODO switch to worker compression to avoid bundling...
-
-// import brotli from 'brotli'; - brotli has problems with decompress in browsers
-// import brotliDecompress from 'brotli/decompress';
-import lz4js from 'lz4js';
-// import lzo from 'lzo';
-// import {ZstdCodec} from 'zstd-codec';
-
-// Inject large dependencies through Compression constructor options
-const modules = {
-  // brotli has problems with decompress in browsers
-  // brotli: {
-  //   decompress: brotliDecompress,
-  //   compress: () => {
-  //     throw new Error('brotli compress');
-  //   }
-  // },
-  lz4js
-  // lzo
-  // 'zstd-codec': ZstdCodec
-};
+import {toArrayBuffer, toUint8Array} from './utils/binary-utils';
 
 /**
  * See https://github.com/apache/parquet-format/blob/master/Compression.md
  */
-// @ts-expect-error
-export const PARQUET_COMPRESSION_METHODS: Record<ParquetCompression, Compression> = {
-  UNCOMPRESSED: new NoCompression(),
-  GZIP: new GZipCompression(),
-  SNAPPY: new SnappyCompression(),
-  BROTLI: new BrotliCompression({modules}),
-  // TODO: Understand difference between LZ4 and LZ4_RAW
-  LZ4: new LZ4Compression({modules}),
-  LZ4_RAW: new LZ4Compression({modules}),
-  //
-  // LZO: new LZOCompression({modules}),
-  ZSTD: new ZstdCompression({modules})
+export const PARQUET_COMPRESSION_METHODS: Partial<Record<ParquetCompression, true>> = {
+  UNCOMPRESSED: true,
+  GZIP: true,
+  SNAPPY: true,
+  BROTLI: true,
+  // TODO: Understand difference between LZ4 and LZ4_RAW.
+  LZ4: true,
+  LZ4_RAW: true,
+  ZSTD: true
 };
 
+/** Native formats available to asynchronous Parquet page decompression. */
+const PARQUET_NATIVE_DECOMPRESSION_FORMATS: Partial<
+  Record<ParquetCompression, NativeDecompressionFormat>
+> = {
+  GZIP: 'gzip',
+  BROTLI: 'brotli',
+  ZSTD: 'zstd'
+};
+
+/** Lazily constructed codec-backed compression implementations. */
+const compressionPromises: Partial<Record<ParquetCompression, Promise<Compression>>> = {};
+
 /**
- * Register compressions that have big external libraries
- * @param options.modules External library dependencies
+ * Registers optional codec modules without eagerly loading codec-backed implementations.
+ *
+ * @param options.modules External library dependencies.
  */
 export async function preloadCompressions(options?: {modules?: {[key: string]: any}}) {
   registerJSModules(options?.modules);
-  const compressions = Object.values(PARQUET_COMPRESSION_METHODS);
-  return await Promise.all(
-    compressions.map((compression) => compression.preload(options?.modules))
-  );
 }
 
 /**
  * Deflate a value using compression method `method`
  */
-export async function deflate(method: ParquetCompression, value: Buffer): Promise<Buffer> {
-  const compression = PARQUET_COMPRESSION_METHODS[method];
-  if (!compression) {
+export async function deflate(method: ParquetCompression, value: Uint8Array): Promise<Uint8Array> {
+  if (!(method in PARQUET_COMPRESSION_METHODS)) {
     throw new Error(`parquet: invalid compression method: ${method}`);
   }
+  if (method === 'UNCOMPRESSED') {
+    return value;
+  }
+
+  const compression = await getParquetCompression(method);
   const inputArrayBuffer = toArrayBuffer(value);
   const compressedArrayBuffer = await compression.compress(inputArrayBuffer);
-  return toBuffer(compressedArrayBuffer);
+  return toUint8Array(compressedArrayBuffer);
 }
 
 /**
@@ -102,106 +72,91 @@ export async function deflate(method: ParquetCompression, value: Buffer): Promis
  */
 export async function decompress(
   method: ParquetCompression,
-  value: Buffer,
+  value: Uint8Array,
   size: number
-): Promise<Buffer> {
-  const compression = PARQUET_COMPRESSION_METHODS[method];
-  if (!compression) {
+): Promise<Uint8Array> {
+  if (!(method in PARQUET_COMPRESSION_METHODS)) {
     throw new Error(`parquet: invalid compression method: ${method}`);
   }
   const inputArrayBuffer = toArrayBuffer(value);
-  const compressedArrayBuffer = await compression.decompress(inputArrayBuffer, size);
-  return toBuffer(compressedArrayBuffer);
-}
-
-/*
- * Inflate a value using compression method `method`
- */
-export function inflate(method: ParquetCompression, value: Buffer, size: number): Buffer {
-  if (!(method in PARQUET_COMPRESSION_METHODS)) {
-    throw new Error(`invalid compression method: ${method}`);
+  if (method === 'UNCOMPRESSED') {
+    return toUint8Array(inputArrayBuffer);
   }
-  // @ts-ignore
-  return PARQUET_COMPRESSION_METHODS[method].inflate(value, size);
-}
 
-/*
-function deflate_identity(value: Buffer): Buffer {
-  return value;
-}
-
-function deflate_gzip(value: Buffer): Buffer {
-  return zlib.gzipSync(value);
-}
-
-function deflate_snappy(value: Buffer): Buffer {
-  return snappyjs.compress(value);
-}
-
-function deflate_lzo(value: Buffer): Buffer {
-  lzo = lzo || Util.load('lzo');
-  return lzo.compress(value);
-}
-
-function deflate_brotli(value: Buffer): Buffer {
-  brotli = brotli || Util.load('brotli');
-  const result = brotli.compress(value, {
-    mode: 0,
-    quality: 8,
-    lgwin: 22
-  });
-  return result ? Buffer.from(result) : Buffer.alloc(0);
-}
-
-function deflate_lz4(value: Buffer): Buffer {
-  lz4js = lz4js || Util.load('lz4js');
-  try {
-    // let result = Buffer.alloc(lz4js.encodeBound(value.length));
-    // const compressedSize = lz4.encodeBlock(value, result);
-    // // remove unnecessary bytes
-    // result = result.slice(0, compressedSize);
-    // return result;
-    return Buffer.from(lz4js.compress(value));
-  } catch (err) {
-    throw err;
+  const nativeFormat = PARQUET_NATIVE_DECOMPRESSION_FORMATS[method];
+  if (nativeFormat && shouldUseNativeDecompressionStream(method)) {
+    const nativeOutput = await decompressWithNativeDecompressionStream(
+      inputArrayBuffer,
+      nativeFormat
+    );
+    if (nativeOutput) {
+      return toUint8Array(nativeOutput);
+    }
   }
-}
-function inflate_identity(value: Buffer): Buffer {
-  return value;
-}
 
-function inflate_gzip(value: Buffer): Buffer {
-  return zlib.gunzipSync(value);
-}
-
-function inflate_snappy(value: Buffer): Buffer {
-  return snappyjs.uncompress(value);
+  const compression = await getParquetCompression(method);
+  const compressedArrayBuffer = await compression.decompress(inputArrayBuffer, size);
+  return toUint8Array(compressedArrayBuffer);
 }
 
-function inflate_lzo(value: Buffer, size: number): Buffer {
-  lzo = lzo || Util.load('lzo');
-  return lzo.decompress(value, size);
+/**
+ * Returns whether a native stream can take precedence for one Parquet compression method.
+ *
+ * @param method Parquet compression method.
+ * @returns Whether no explicitly registered codec should take precedence.
+ */
+function shouldUseNativeDecompressionStream(method: ParquetCompression): boolean {
+  if (method === 'BROTLI') {
+    return !getJSModuleOrNull('brotli');
+  }
+  if (method === 'ZSTD') {
+    return !getJSModuleOrNull('zstd-codec');
+  }
+  return true;
 }
 
-function inflate_lz4(value: Buffer, size: number): Buffer {
-  lz4js = lz4js || Util.load('lz4js');
-  try {
-    // let result = Buffer.alloc(size);
-    // const uncompressedSize = lz4js.decodeBlock(value, result);
-    // // remove unnecessary bytes
-    // result = result.slice(0, uncompressedSize);
-    // return result;
-    return Buffer.from(lz4js.decompress(value, size));
-  } catch (err) {
-    throw err;
-  }
+/**
+ * Loads one codec-backed implementation only after native decompression is unavailable.
+ *
+ * @param method Parquet compression method.
+ * @returns Codec-backed compression implementation.
+ */
+async function getParquetCompression(method: ParquetCompression): Promise<Compression> {
+  compressionPromises[method] ||= createParquetCompression(method);
+  return await compressionPromises[method];
 }
 
-function inflate_brotli(value: Buffer): Buffer {
-  brotli = brotli || Util.load('brotli');
-  if (!value.length) {
-    return Buffer.alloc(0);
+/**
+ * Creates one lazily loaded codec-backed Parquet compression implementation.
+ *
+ * @param method Parquet compression method.
+ * @returns Codec-backed compression implementation.
+ */
+async function createParquetCompression(method: ParquetCompression): Promise<Compression> {
+  switch (method) {
+    case 'GZIP': {
+      const {GZipCompression} = await import('@loaders.gl/compression/gzip-compression');
+      return new GZipCompression();
+    }
+    case 'SNAPPY': {
+      const {SnappyCompression} = await import('@loaders.gl/compression/snappy-compression');
+      return new SnappyCompression();
+    }
+    case 'BROTLI': {
+      const {BrotliCompression} = await import('@loaders.gl/compression/brotli-compression');
+      return new BrotliCompression();
+    }
+    case 'LZ4':
+    case 'LZ4_RAW': {
+      const {LZ4Compression} = await import('@loaders.gl/compression/lz4-compression');
+      const lz4js = getJSModuleOrNull('lz4js') || (await import('lz4js')).default;
+      return new LZ4Compression({modules: {lz4js}});
+    }
+    case 'ZSTD': {
+      const {ZstdCompression} = await import('@loaders.gl/compression/zstd-compression');
+      return new ZstdCompression();
+    }
+    default:
+      throw new Error(`parquet: invalid compression method: ${method}`);
   }
-  return Buffer.from(brotli.decompress(value));
 }
-*/

package/src/parquetjs/encoder/dictionary-planner.ts

@@ -0,0 +1,99 @@
+// loaders.gl
+// SPDX-License-Identifier: MIT
+// Copyright (c) vis.gl contributors
+
+import {PARQUET_CODECS} from '../codecs/index';
+import type {ParquetField} from '../schema/declare';
+
+/** Dictionary selection policy for one Parquet column chunk. */
+export type ParquetDictionaryPolicy = boolean | 'auto';
+
+/** A chunk-wide dictionary and the index stream replacing its non-null values. */
+export interface ParquetDictionaryPlan {
+  /** Unique physical values written to the dictionary page. */
+  readonly values: unknown[];
+  /** Dictionary index corresponding to every non-null column value. */
+  readonly indices: number[];
+  /** PLAIN-encoded dictionary page payload. */
+  readonly encodedValues: Uint8Array;
+  /** Bit width used by RLE_DICTIONARY index pages. */
+  readonly bitWidth: number;
+}
+
+/** Builds a forced or size-beneficial dictionary plan for a complete column chunk. */
+export function planDictionary(
+  column: ParquetField,
+  values: readonly unknown[],
+  policy: ParquetDictionaryPolicy,
+  dictionaryPageSizeLimit: number
+): ParquetDictionaryPlan | undefined {
+  if (!Number.isInteger(dictionaryPageSizeLimit) || dictionaryPageSizeLimit <= 0) {
+    throw new Error(
+      `Parquet dictionary page size limit must be a positive integer, received ${dictionaryPageSizeLimit}`
+    );
+  }
+  if (policy === false || values.length === 0 || isDictionaryEncoding(column.encoding!)) {
+    return undefined;
+  }
+
+  const dictionaryValues: unknown[] = [];
+  const dictionaryIndices = new Map<string, number>();
+  const indices = new Array<number>(values.length);
+  for (let valueIndex = 0; valueIndex < values.length; valueIndex++) {
+    const value = values[valueIndex];
+    const key = getDictionaryKey(value);
+    let dictionaryIndex = dictionaryIndices.get(key);
+    if (dictionaryIndex === undefined) {
+      dictionaryIndex = dictionaryValues.length;
+      dictionaryIndices.set(key, dictionaryIndex);
+      dictionaryValues.push(value);
+    }
+    indices[valueIndex] = dictionaryIndex;
+  }
+
+  const encodedValues = PARQUET_CODECS.PLAIN.encodeValues(column.primitiveType!, dictionaryValues, {
+    typeLength: column.typeLength
+  });
+  if (encodedValues.length > dictionaryPageSizeLimit) {
+    return undefined;
+  }
+  const bitWidth = dictionaryValues.length <= 1 ? 0 : Math.ceil(Math.log2(dictionaryValues.length));
+  if (policy === 'auto') {
+    const encodedIndices = PARQUET_CODECS.RLE_DICTIONARY.encodeValues(
+      column.primitiveType!,
+      indices,
+      {bitWidth}
+    );
+    const encodedPrimaryValues = PARQUET_CODECS[column.encoding!].encodeValues(
+      column.primitiveType!,
+      values as unknown[],
+      {typeLength: column.typeLength, bitWidth: column.typeLength}
+    );
+    if (encodedValues.length + encodedIndices.length >= encodedPrimaryValues.length) {
+      return undefined;
+    }
+  }
+
+  return {values: dictionaryValues, indices, encodedValues, bitWidth};
+}
+
+/** Returns a stable content key for one physical Parquet value. */
+function getDictionaryKey(value: unknown): string {
+  if (value instanceof Uint8Array) {
+    let key = `bytes:${value.length}:`;
+    for (const byte of value) {
+      key += String.fromCharCode(byte);
+    }
+    return key;
+  }
+  if (typeof value === 'number') {
+    if (Number.isNaN(value)) return 'number:NaN';
+    if (Object.is(value, -0)) return 'number:-0';
+  }
+  return `${typeof value}:${String(value)}`;
+}
+
+/** Returns whether the selected primary encoding already requires a dictionary page. */
+function isDictionaryEncoding(encoding: string): boolean {
+  return encoding === 'PLAIN_DICTIONARY' || encoding === 'RLE_DICTIONARY';
+}

package/src/parquetjs/encoder/page-planner.ts

@@ -0,0 +1,83 @@
+// loaders.gl
+// SPDX-License-Identifier: MIT
+// Copyright (c) vis.gl contributors
+
+import type {ParquetColumnChunk, ParquetField} from '../schema/declare';
+
+/** One row-aligned slice of a shredded column prepared for page encoding. */
+export interface PlannedParquetPage {
+  /** Shredded levels and non-null values contained in this page. */
+  readonly data: ParquetColumnChunk;
+  /** Number of top-level rows represented by the page. */
+  readonly rowCount: number;
+}
+
+/** Splits a shredded column into target-sized pages without splitting a top-level row. */
+export function planColumnPages(
+  column: ParquetField,
+  data: ParquetColumnChunk,
+  rowCount: number,
+  pageSize: number
+): PlannedParquetPage[] {
+  if (!Number.isInteger(pageSize) || pageSize <= 0) {
+    throw new Error(`Parquet page size must be a positive integer, received ${pageSize}`);
+  }
+  if (data.rlevels.length < data.count || data.dlevels.length < data.count) {
+    throw new Error(`Invalid shredded column ${column.key}: level buffers are shorter than count`);
+  }
+  if (rowCount === 0) {
+    return [];
+  }
+
+  const rowStarts = [0];
+  for (let levelIndex = 1; levelIndex < data.count; levelIndex++) {
+    if (data.rlevels[levelIndex] === 0) {
+      rowStarts.push(levelIndex);
+    }
+  }
+  rowStarts.push(data.count);
+  if (rowStarts.length !== rowCount + 1) {
+    throw new Error(
+      `Invalid shredded column ${column.key}: found ${rowStarts.length - 1} rows, expected ${rowCount}`
+    );
+  }
+
+  const pages: PlannedParquetPage[] = [];
+  let pageStartRow = 0;
+  let valueStart = 0;
+  while (pageStartRow < rowCount) {
+    const levelStart = rowStarts[pageStartRow];
+    let pageEndRow = pageStartRow + 1;
+    while (pageEndRow < rowCount && rowStarts[pageEndRow + 1] - levelStart <= pageSize) {
+      pageEndRow++;
+    }
+
+    const levelEnd = rowStarts[pageEndRow];
+    let valueCount = 0;
+    for (let levelIndex = levelStart; levelIndex < levelEnd; levelIndex++) {
+      if (data.dlevels[levelIndex] === column.dLevelMax) {
+        valueCount++;
+      }
+    }
+    const valueEnd = valueStart + valueCount;
+    pages.push({
+      rowCount: pageEndRow - pageStartRow,
+      data: {
+        rlevels: data.rlevels.slice(levelStart, levelEnd),
+        dlevels: data.dlevels.slice(levelStart, levelEnd),
+        values: data.values.slice(valueStart, valueEnd),
+        count: levelEnd - levelStart,
+        pageHeaders: []
+      }
+    });
+    pageStartRow = pageEndRow;
+    valueStart = valueEnd;
+  }
+
+  if (valueStart !== data.values.length) {
+    throw new Error(
+      `Invalid shredded column ${column.key}: consumed ${valueStart} values, expected ${data.values.length}`
+    );
+  }
+  return pages;
+}