benchmark.js

import { createWriteStream, promises as fs } from 'fs'
import { asyncBufferFromFile } from 'hyparquet'
import { fileWriter } from 'hyparquet-writer'
import { pipeline } from 'stream/promises'
import { createIndex } from './src/createIndex.js'
import { parquetFind } from './src/parquetFind.js'

const url = 'https://huggingface.co/datasets/wikimedia/wikipedia/resolve/main/20231101.en/train-00000-of-00041.parquet'
const filename = 'example.parquet'
const indexFilename = 'example.index.parquet'

// Download source if needed
let stat = await fs.stat(filename).catch(() => undefined)
if (!stat) {
  console.log('downloading ' + url)
  const res = await fetch(url)
  if (!res.ok) throw new Error(res.statusText)
  await pipeline(res.body, createWriteStream(filename))
  stat = await fs.stat(filename)
}

// Build index if needed
let indexStat = await fs.stat(indexFilename).catch(() => undefined)
if (!indexStat) {
  console.log('building ' + indexFilename)
  const t0 = performance.now()
  const sourceFile = await asyncBufferFromFile(filename)
  const indexFile = fileWriter(indexFilename)
  await createIndex({ sourceFile, indexFile })
  indexStat = await fs.stat(indexFilename)
  console.log(`built in ${((performance.now() - t0) / 1000).toFixed(1)}s`)
}

console.log(`source: ${(stat.size / 1024 / 1024).toFixed(1)} MB`)
console.log(`index:  ${(indexStat.size / 1024 / 1024).toFixed(1)} MB (${(indexStat.size / stat.size * 100).toFixed(2)}% of source)`)

/**
 * @param {import('hyparquet').AsyncBuffer} buf
 * @returns {import('hyparquet').AsyncBuffer & {fetches: number, bytes: number}}
 */
function instrument(buf) {
  const wrapper = {
    byteLength: buf.byteLength,
    fetches: 0,
    bytes: 0,
    /**
     * @param {number} start
     * @param {number} [end]
     * @returns {Promise<ArrayBuffer> | ArrayBuffer}
     */
    slice(start, end) {
      wrapper.fetches += 1
      wrapper.bytes += (end ?? buf.byteLength) - start
      return buf.slice(start, end)
    },
  }
  return wrapper
}

const queries = [
  'eigenvalue',
  'petrichor',
  'serverless',
  'quantum entanglement',
  'wikipedia',
]

for (const limit of [Infinity, 10]) {
  console.log()
  console.log('limit = ' + (limit === Infinity ? 'all' : limit))
  console.log('query                 matches  ms     idx_KB  src_MB  total')
  console.log('--------------------- -------  -----  ------  ------  -----')
  for (const query of queries) {
    const idx = instrument(await asyncBufferFromFile(indexFilename))
    const src = instrument(await asyncBufferFromFile(filename))
    const t0 = performance.now()
    let matches = 0
    // eslint-disable-next-line no-unused-vars
    for await (const _row of parquetFind({
      url: filename,
      sourceFile: src,
      indexFile: idx,
      query,
      limit,
    })) matches += 1
    const ms = performance.now() - t0
    const total = (idx.bytes + src.bytes) / 1024 / 1024
    console.log(
      query.padEnd(21) +
      ' ' + String(matches).padStart(7) +
      '  ' + ms.toFixed(0).padStart(5) +
      '  ' + (idx.bytes / 1024).toFixed(0).padStart(6) +
      '  ' + (src.bytes / 1024 / 1024).toFixed(0).padStart(6) +
      '  ' + total.toFixed(1).padStart(5) + ' MB'
    )
  }
}