Line data    Source code

       1              : import type * as arrow from 'apache-arrow';
       2              : 
       3              : import type { EmbeddedComputeDatabase, EmbeddedExternalDirectoryDatabase, EmbeddedExternalFileDatabase } from '../database/embedded_database.js';
       4              : import { getOPFSRoot } from '../opfs_root.js';
       5              : 
       6              : /// Hyper writes Parquet exports below this OPFS directory, they are read back and removed right after.
       7            1 : export const PARQUET_EXPORT_DIRECTORY = 'exports';
       8            1 : const PARQUET_EXPORT_FILE = 'result.parquet';
       9              : 
      10              : export type ParquetExportDatabase = EmbeddedComputeDatabase & EmbeddedExternalDirectoryDatabase;
      11              : 
      12            1 : export function supportsParquetExport(database: EmbeddedComputeDatabase | null): database is ParquetExportDatabase {
      13            1 :     return database != null && 'registerExternalDirectory' in database && 'removeExternalDirectory' in database;
      14            1 : }
      15              : 
      16              : /// Write an Arrow table as Parquet file to a path below a directory that is mounted read-write
      17            1 : export async function copyArrowTableAsParquet(database: EmbeddedComputeDatabase, table: arrow.Table, path: string): Promise<void> {
      18            1 :     const tableName = `__parquet_copy_${crypto.randomUUID().replace(/-/g, '_')}`;
      19            1 :     const connection = await database.connect({ pooled: true });
      20            1 :     try {
      21            1 :         await connection.insertArrowTable(table, { name: tableName });
      22            1 :         try {
      23            1 :             await connection.execute(`COPY (SELECT * FROM "${tableName}") TO '${path}' WITH (format => 'parquet')`);
      24            1 :         } finally {
      25            1 :             await connection.execute(`DROP TABLE IF EXISTS "${tableName}"`);
      26            1 :         }
      27            1 :     } finally {
      28            1 :         await connection.close();
      29            1 :     }
      30            1 : }
      31              : 
      32              : /// Write an Arrow table as Parquet file to OPFS with Hyper and read the file back.
      33              : /// Hyper only copies into registered targets, so every export mounts its own OPFS directory read-write.
      34            1 : export async function exportArrowTableAsParquet(
      35            1 :     database: ParquetExportDatabase,
      36            1 :     table: arrow.Table,
      37            1 :     getRoot: () => Promise<FileSystemDirectoryHandle> = getOPFSRoot,
      38            1 : ): Promise<Uint8Array> {
      39            1 :     const root = await getRoot();
      40            1 :     const exports = await root.getDirectoryHandle(PARQUET_EXPORT_DIRECTORY, { create: true });
      41            1 :     const id = crypto.randomUUID();
      42            1 :     const directory = await exports.getDirectoryHandle(id, { create: true });
      43            1 :     try {
      44            1 :         const path = await database.registerExternalDirectory(directory, `/mnt/external/parquet-export-${id}`, 'readwrite');
      45            1 :         try {
      46            1 :             await copyArrowTableAsParquet(database, table, `${path}/${PARQUET_EXPORT_FILE}`);
      47            1 :         } finally {
      48            1 :             await database.removeExternalDirectory(path);
      49            1 :         }
      50            1 :         const file = await (await directory.getFileHandle(PARQUET_EXPORT_FILE)).getFile();
      51            1 :         return new Uint8Array(await file.arrayBuffer());
      52            1 :     } finally {
      53            1 :         await exports.removeEntry(id, { recursive: true });
      54            1 :     }
      55            1 : }
      56              : 
      57              : /// A database that writes and reads Parquet files: it mounts directories to write them and files to read them, and
      58              : /// refreshes directory mounts whose files were removed
      59              : export type ParquetDatabase = ParquetExportDatabase & EmbeddedExternalFileDatabase;
      60              : 
      61            1 : export function supportsParquet(database: EmbeddedComputeDatabase | null): database is ParquetDatabase {
      62            1 :     return supportsParquetExport(database) && 'registerExternalFile' in database && 'removeExternalFile' in database
      63            1 :         && 'refreshExternalDirectory' in database;
      64            1 : }