Stream index creation through byte arrays

benchmark/build-wildchat-default.js

@@ -83,9 +83,11 @@ try {
 }
 
 const outputStat = await fs.stat(outputPath)
+const resourceUsage = process.resourceUsage()
 console.log(JSON.stringify({
   complete: true,
   elapsedMinutes: (performance.now() - started) / 60000,
+  peakRssGB: resourceUsage.maxRSS / 1e6,
   outputPath,
   outputBytes: outputStat.size,
 }))

package.json

@@ -44,7 +44,7 @@
   "dependencies": {
     "hyparquet": "1.27.1",
     "hyparquet-compressors": "1.1.1",
-    "hyparquet-writer": "0.16.3"
+    "hyparquet-writer": "0.16.4"
   },
   "devDependencies": {
     "@eslint/js": "10.0.1",

src/createIndex.js

@@ -1,7 +1,8 @@
-import { parquetMetadataAsync, parquetReadObjects } from 'hyparquet'
+import { parquetMetadataAsync, parquetRead } from 'hyparquet'
+import { DEFAULT_PARSERS } from 'hyparquet/src/convert.js'
 import { ParquetWriter, schemaFromColumnData } from 'hyparquet-writer'
 import { chooseIndexRowGroupSize, defaultBlockSize, defaultHexShapeNgramLength, defaultIndexPageSize, defaultNgramChars, defaultNgramLength, defaultShapeChars, defaultShapeNgramLength, defaultShortShapeNgramLength, hypGrepVersion } from './constants.js'
-import { decodeIndexToken, extractIndexTokenCodes, supportsNumericIndexTokens } from './indexTokens.js'
+import { decodeIndexTokenBytes, extractIndexTokenBytes, extractIndexTokenCodes, supportsNumericIndexTokens } from './indexTokens.js'
 import { extractNgrams } from './ngrams.js'
 import { NumericPostingBuffer, postingTokenRadix } from './postingBuffer.js'
 import { extractStructuralNgrams } from './shape.js'
@@ -64,6 +65,20 @@ export async function createIndex({
     textColumns = requestedColumns
   }
   const useNumericTokens = supportsNumericIndexTokens(ngramLength, ngramChars)
+  // Hyparquet merges parser overrides with its defaults. Preserve annotated
+  // STRING values as their decoded Parquet byte arrays on the numeric path so
+  // token extraction can scan UTF-8 directly without materializing JS strings.
+  const byteStringParser = {
+    ...DEFAULT_PARSERS,
+    /**
+     * @param {Uint8Array} bytes
+     * @returns {Uint8Array}
+     */
+    stringFromBytes(bytes) {
+      return bytes
+    },
+  }
+  const sourceParserOptions = useNumericTokens ? { parsers: byteStringParser } : {}
 
   // Postings partitioned by n-gram prefix: prefix -> (n-gram -> blockIds, in
   // ascending order). A single Map keyed on every distinct n-gram overflows
@@ -151,28 +166,58 @@ export async function createIndex({
   let groupStart = 0
   for (const group of metadata.row_groups) {
     const groupRows = Number(group.num_rows)
-    const rows = await parquetReadObjects({
+    /** @type {Map<string, import('hyparquet').ColumnData[]>} */
+    const chunksByColumn = new Map(textColumns.map(name => [name, []]))
+    await parquetRead({
       file: sourceFile,
       metadata,
       rowStart: groupStart,
       rowEnd: groupStart + groupRows,
       columns: textColumns,
+      ...sourceParserOptions,
+      onChunk(chunk) {
+        const chunks = chunksByColumn.get(chunk.columnName)
+        if (chunks) chunks.push(chunk)
+      },
     })
-    for (const row of rows) {
-      if (useNumericTokens) {
-        collectRowTokenCodes(row, textColumns, blockTokenCodes)
-      } else {
-        collectRowNgrams(
-          row,
-          textColumns,
-          ngramLength,
-          ngramChars,
-          shapeNgramLength,
-          shortShapeNgramLength,
-          shapeChars,
-          hexShapeNgramLength,
-          blockNgrams
-        )
+
+    // Walk the decoded columns in place. parquetReadObjects transposes these
+    // chunks into one temporary object per source row, which is unnecessary
+    // when index creation only needs to visit every selected string value.
+    const columnCursors = textColumns.map(name => ({
+      chunks: chunksByColumn.get(name)?.sort((a, b) => a.rowStart - b.rowStart) ?? [],
+      index: 0,
+    }))
+    for (let row = groupStart; row < groupStart + groupRows; row += 1) {
+      for (const cursor of columnCursors) {
+        while (cursor.index < cursor.chunks.length && row >= cursor.chunks[cursor.index].rowEnd) {
+          cursor.index += 1
+        }
+        const chunk = cursor.chunks[cursor.index]
+        if (!chunk || row < chunk.rowStart) {
+          throw new Error(`Missing decoded text column data at source row ${row}`)
+        }
+        const value = chunk.columnData[row - chunk.rowStart]
+        if (useNumericTokens) {
+          if (value instanceof Uint8Array) {
+            extractIndexTokenBytes(value, blockTokenCodes)
+          } else if (typeof value === 'string') {
+            // A custom source parser may already have converted the value.
+            extractIndexTokenCodes(value, blockTokenCodes)
+          }
+        } else {
+          if (typeof value !== 'string') continue
+          collectTextNgrams(
+            value,
+            ngramLength,
+            ngramChars,
+            shapeNgramLength,
+            shortShapeNgramLength,
+            shapeChars,
+            hexShapeNgramLength,
+            blockNgrams
+          )
+        }
       }
       if (++rowsInBlock === blockSize) flushBlock()
     }
@@ -219,7 +264,7 @@ export async function createIndex({
   const writer = new ParquetWriter({ writer: indexFile, schema, kvMetadata })
 
   /**
-   * @param {string[]} ngramData
+   * @param {(string | Uint8Array)[]} ngramData
    * @param {number[]} blockData
    */
   async function writeRowGroup(ngramData, blockData) {
@@ -248,7 +293,7 @@ export async function createIndex({
     })
   }
 
-  /** @type {string[]} */
+  /** @type {(string | Uint8Array)[]} */
   const ngramBatch = []
   /** @type {number[]} */
   const blockBatch = []
@@ -264,7 +309,7 @@ export async function createIndex({
       for (let suffix = 0; suffix < tokenEnds.length; suffix += 1) {
         const postingEnd = tokenEnds[suffix]
         if (postingEnd === postingStart) continue
-        const ngram = decodeIndexToken(prefix * postingTokenRadix + suffix)
+        const ngram = decodeIndexTokenBytes(prefix * postingTokenRadix + suffix)
         for (let i = postingStart; i < postingEnd; i += 1) {
           ngramBatch.push(ngram)
           blockBatch.push(sortedBlockIds[i])
@@ -311,28 +356,11 @@ export async function createIndex({
 }
 
 /**
- * Collect compact numeric tokens for the default index configuration.
- *
- * @param {Record<string, any>} row
- * @param {string[]} textColumns
- * @param {Set<number>} tokens block accumulator to add into
- * @returns {void}
- */
-function collectRowTokenCodes(row, textColumns, tokens) {
-  if (!row) return
-  for (const columnName of textColumns) {
-    const value = row[columnName]
-    if (typeof value === 'string') extractIndexTokenCodes(value, tokens)
-  }
-}
-
-/**
- * Collect the distinct n-grams of one row into a block's accumulator set:
+ * Collect the distinct n-grams of one text value into a block's accumulator:
  * ordinary character n-grams plus shape (character-class skeleton) n-grams,
  * which share the column but are namespaced by SHAPE_PREFIX.
  *
- * @param {Record<string, any>} row
- * @param {string[]} textColumns
+ * @param {string} value
  * @param {number} n
  * @param {string} chars extra characters kept inside n-gram runs
  * @param {number} shapeN primary shape n-gram length
@@ -341,9 +369,8 @@ function collectRowTokenCodes(row, textColumns, tokens) {
  * @param {number} hexShapeN hexadecimal shape n-gram length
  * @param {Set<string>} ngrams block accumulator to add into
  */
-function collectRowNgrams(
-  row,
-  textColumns,
+function collectTextNgrams(
+  value,
   n,
   chars,
   shapeN,
@@ -352,21 +379,16 @@ function collectRowNgrams(
   hexShapeN,
   ngrams
 ) {
-  if (!row) return
-  for (const columnName of textColumns) {
-    const value = row[columnName]
-    if (typeof value !== 'string') continue
-    if (value.length >= n) {
-      for (const g of extractNgrams(value, n, chars)) ngrams.add(g)
-    }
-    // All structural layers share one source traversal and stream directly
-    // into the block accumulator.
-    extractStructuralNgrams(
-      value,
-      [shapeN, shortShapeN],
-      hexShapeN,
-      shapeChars,
-      ngrams
-    )
+  if (value.length >= n) {
+    for (const g of extractNgrams(value, n, chars)) ngrams.add(g)
   }
+  // All structural layers share one source traversal and stream directly
+  // into the block accumulator.
+  extractStructuralNgrams(
+    value,
+    [shapeN, shortShapeN],
+    hexShapeN,
+    shapeChars,
+    ngrams
+  )
 }