Use numeric tokens for faster index builds

package.json

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

src/createIndex.js

@@ -1,6 +1,7 @@
 import { parquetMetadataAsync, parquetReadObjects } from 'hyparquet'
 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 { extractNgrams } from './ngrams.js'
 import { extractStructuralNgrams } from './shape.js'
 import { assertNonNegativeSafeInteger, assertPositiveSafeInteger, getTextColumnsFromSchema } from './utils.js'
@@ -61,6 +62,7 @@ export async function createIndex({
     }
     textColumns = requestedColumns
   }
+  const useNumericTokens = supportsNumericIndexTokens(ngramLength, ngramChars)
 
   // Postings partitioned by n-gram prefix: prefix -> (n-gram -> blockIds, in
   // ascending order). A single Map keyed on every distinct n-gram overflows
@@ -81,46 +83,69 @@ export async function createIndex({
   const ASCII_CODE_UNITS = 128
   const ASCII_PREFIX_RADIX = ASCII_CODE_UNITS + 1
   const UTF16_PREFIX_RADIX = 65536 + 1
+  const NUMERIC_PREFIX_RADIX = 65536
   /** @type {Array<Map<string, number[]> | undefined>} */
   const asciiBuckets = []
   /** @type {Map<number, Map<string, number[]>>} */
   const unicodeBuckets = new Map()
   /** @type {{ascii: boolean, first: number, id: number, second: number}[]} */
   const activePrefixes = []
+  /** @type {Array<Map<number, number[]> | undefined>} */
+  const numericBuckets = []
+  /** @type {number[]} */
+  const activeNumericPrefixes = []
 
   let blockId = 0
   let rowsInBlock = 0
   let postingRows = 0
   /** @type {Set<string>} */
   let blockNgrams = new Set()
+  /** @type {Set<number>} */
+  let blockTokenCodes = new Set()
 
   /** Move the accumulated block n-grams into the postings buckets. */
   function flushBlock() {
-    postingRows += blockNgrams.size
+    postingRows += useNumericTokens ? blockTokenCodes.size : blockNgrams.size
     if (!Number.isSafeInteger(postingRows)) {
       throw new Error('Index posting row count exceeds the safe integer range')
     }
-    for (const ngram of blockNgrams) {
-      const first = ngram.charCodeAt(0)
-      const second = ngram.length > 1 ? ngram.charCodeAt(1) : -1
-      const ascii = first < ASCII_CODE_UNITS && second < ASCII_CODE_UNITS
-      const id = ascii
-        ? first * ASCII_PREFIX_RADIX + second + 1
-        : first * UTF16_PREFIX_RADIX + second + 1
-      let postings = ascii ? asciiBuckets[id] : unicodeBuckets.get(id)
-      if (!postings) {
-        postings = new Map()
-        if (ascii) asciiBuckets[id] = postings
-        else unicodeBuckets.set(id, postings)
-        activePrefixes.push({ ascii, first, id, second })
+    if (useNumericTokens) {
+      for (const token of blockTokenCodes) {
+        const prefix = Math.floor(token / NUMERIC_PREFIX_RADIX)
+        let postings = numericBuckets[prefix]
+        if (!postings) {
+          postings = new Map()
+          numericBuckets[prefix] = postings
+          activeNumericPrefixes.push(prefix)
+        }
+        const existing = postings.get(token)
+        if (existing) existing.push(blockId)
+        else postings.set(token, [blockId])
+      }
+    } else {
+      for (const ngram of blockNgrams) {
+        const first = ngram.charCodeAt(0)
+        const second = ngram.length > 1 ? ngram.charCodeAt(1) : -1
+        const ascii = first < ASCII_CODE_UNITS && second < ASCII_CODE_UNITS
+        const id = ascii
+          ? first * ASCII_PREFIX_RADIX + second + 1
+          : first * UTF16_PREFIX_RADIX + second + 1
+        let postings = ascii ? asciiBuckets[id] : unicodeBuckets.get(id)
+        if (!postings) {
+          postings = new Map()
+          if (ascii) asciiBuckets[id] = postings
+          else unicodeBuckets.set(id, postings)
+          activePrefixes.push({ ascii, first, id, second })
+        }
+        const existing = postings.get(ngram)
+        if (existing) existing.push(blockId)
+        else postings.set(ngram, [blockId])
       }
-      const existing = postings.get(ngram)
-      if (existing) existing.push(blockId)
-      else postings.set(ngram, [blockId])
     }
     blockId += 1
     rowsInBlock = 0
     blockNgrams = new Set()
+    blockTokenCodes = new Set()
   }
 
   // Read the source once, one row group per read, and slice blocks in
@@ -139,17 +164,21 @@ export async function createIndex({
       columns: textColumns,
     })
     for (const row of rows) {
-      collectRowNgrams(
-        row,
-        textColumns,
-        ngramLength,
-        ngramChars,
-        shapeNgramLength,
-        shortShapeNgramLength,
-        shapeChars,
-        hexShapeNgramLength,
-        blockNgrams
-      )
+      if (useNumericTokens) {
+        collectRowTokenCodes(row, textColumns, blockTokenCodes)
+      } else {
+        collectRowNgrams(
+          row,
+          textColumns,
+          ngramLength,
+          ngramChars,
+          shapeNgramLength,
+          shortShapeNgramLength,
+          shapeChars,
+          hexShapeNgramLength,
+          blockNgrams
+        )
+      }
       if (++rowsInBlock === blockSize) flushBlock()
     }
     groupStart += groupRows
@@ -228,27 +257,52 @@ export async function createIndex({
   const ngramBatch = []
   /** @type {number[]} */
   const blockBatch = []
-  activePrefixes.sort((a, b) => a.first - b.first || a.second - b.second)
-  for (const prefix of activePrefixes) {
-    const postings = prefix.ascii
-      ? asciiBuckets[prefix.id]
-      : unicodeBuckets.get(prefix.id)
-    if (prefix.ascii) asciiBuckets[prefix.id] = undefined
-    else unicodeBuckets.delete(prefix.id)
-    if (!postings) continue
-    for (const ngram of Array.from(postings.keys()).sort()) {
-      const blocks = postings.get(ngram)
-      postings.delete(ngram)
-      if (!blocks) continue
-      for (const id of blocks) {
-        ngramBatch.push(ngram)
-        blockBatch.push(id)
+  if (useNumericTokens) {
+    activeNumericPrefixes.sort((a, b) => a - b)
+    for (const prefix of activeNumericPrefixes) {
+      const postings = numericBuckets[prefix]
+      numericBuckets[prefix] = undefined
+      if (!postings) continue
+      for (const token of Array.from(postings.keys()).sort((a, b) => a - b)) {
+        const blocks = postings.get(token)
+        postings.delete(token)
+        if (!blocks) continue
+        const ngram = decodeIndexToken(token)
+        for (const id of blocks) {
+          ngramBatch.push(ngram)
+          blockBatch.push(id)
+        }
+        if (ngramBatch.length >= resolvedIndexRowGroupSize) {
+          await writeRowGroup(
+            ngramBatch.splice(0, resolvedIndexRowGroupSize),
+            blockBatch.splice(0, resolvedIndexRowGroupSize)
+          )
+        }
       }
-      if (ngramBatch.length >= resolvedIndexRowGroupSize) {
-        await writeRowGroup(
-          ngramBatch.splice(0, resolvedIndexRowGroupSize),
-          blockBatch.splice(0, resolvedIndexRowGroupSize)
-        )
+    }
+  } else {
+    activePrefixes.sort((a, b) => a.first - b.first || a.second - b.second)
+    for (const prefix of activePrefixes) {
+      const postings = prefix.ascii
+        ? asciiBuckets[prefix.id]
+        : unicodeBuckets.get(prefix.id)
+      if (prefix.ascii) asciiBuckets[prefix.id] = undefined
+      else unicodeBuckets.delete(prefix.id)
+      if (!postings) continue
+      for (const ngram of Array.from(postings.keys()).sort()) {
+        const blocks = postings.get(ngram)
+        postings.delete(ngram)
+        if (!blocks) continue
+        for (const id of blocks) {
+          ngramBatch.push(ngram)
+          blockBatch.push(id)
+        }
+        if (ngramBatch.length >= resolvedIndexRowGroupSize) {
+          await writeRowGroup(
+            ngramBatch.splice(0, resolvedIndexRowGroupSize),
+            blockBatch.splice(0, resolvedIndexRowGroupSize)
+          )
+        }
       }
     }
   }
@@ -258,6 +312,22 @@ export async function createIndex({
   await writer.finish()
 }
 
+/**
+ * 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:
  * ordinary character n-grams plus shape (character-class skeleton) n-grams,