diff --git a/benchmark/build-wildchat-default.js b/benchmark/build-wildchat-default.js index 2997f35..e16ea58 100644 --- a/benchmark/build-wildchat-default.js +++ b/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, })) diff --git a/package.json b/package.json index b9c0815..cf6e7d0 100644 --- a/package.json +++ b/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", diff --git a/src/createIndex.js b/src/createIndex.js index b34f60e..14a07cc 100644 --- a/src/createIndex.js +++ b/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} */ + 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} row - * @param {string[]} textColumns - * @param {Set} 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} 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} 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 + ) } diff --git a/src/indexTokens.js b/src/indexTokens.js index d1ca03b..6148fef 100644 --- a/src/indexTokens.js +++ b/src/indexTokens.js @@ -42,6 +42,9 @@ const ordinarySymbols = new Int8Array(128).fill(-1) for (let i = 0; i < ORDINARY_ALPHABET.length; i += 1) { ordinarySymbols[ORDINARY_ALPHABET.charCodeAt(i)] = i } +for (let codeUnit = 65; codeUnit <= 90; codeUnit += 1) { + ordinarySymbols[codeUnit] = ordinarySymbols[codeUnit + 32] +} const shapeSymbols = new Int8Array(128).fill(-1) for (let i = 0; i < SHAPE_ALPHABET.length; i += 1) { @@ -86,6 +89,26 @@ export function extractIndexTokenCodes(text, out = new Set()) { return out } +/** + * Add default index tokens directly from UTF-8 source bytes. + * + * Indexed symbols are ASCII, so non-ASCII code points are boundaries except + * for the two Unicode lowercase mappings that introduce ASCII: U+0130 becomes + * `i` followed by a combining-mark boundary, and U+212A becomes `k`. Handling + * those sequences explicitly makes this byte path match the string tokenizer + * without allocating a decoded copy of every source value. + * + * @param {Uint8Array} bytes UTF-8 source text + * @param {Set} [out] block accumulator + * @returns {Set} + */ +export function extractIndexTokenBytes(bytes, out = new Set()) { + if (!(bytes instanceof Uint8Array)) return out + extractOrdinaryByteCodes(bytes, out) + extractStructuralByteCodes(bytes, out) + return out +} + /** * Decode a numeric build token into the exact string stored in the index. * @@ -121,6 +144,49 @@ export function decodeIndexToken(token) { ) } +/** + * Decode a numeric build token directly into the UTF-8 bytes stored in the + * index. Every token symbol is ASCII, including the structural namespace + * prefixes, so no TextEncoder pass is needed. + * + * @param {number} token numeric index token + * @returns {Uint8Array} + */ +export function decodeIndexTokenBytes(token) { + if (!Number.isSafeInteger(token) || token < 0 || token >= TOKEN_SPACE) { + throw new Error(`Invalid numeric index token: ${token}`) + } + if (token < HEX_OFFSET) { + const digits = decodeDigits(token, SHAPE_LENGTH, SHAPE_BASE) + const bytes = new Uint8Array(SHAPE_LENGTH + 1) + bytes[0] = SHAPE_PREFIX.charCodeAt(0) + let length = 1 + for (const digit of digits) { + if (digit === 0) break + bytes[length++] = SHAPE_ALPHABET.charCodeAt(digit - 1) + } + return bytes.subarray(0, length) + } + if (token < ORDINARY_OFFSET) { + token -= HEX_OFFSET + const bytes = new Uint8Array(HEX_LENGTH + 1) + bytes[0] = HEX_SHAPE_PREFIX.charCodeAt(0) + for (let i = HEX_LENGTH; i > 0; i -= 1) { + bytes[i] = HEX_ALPHABET.charCodeAt(token % HEX_BASE) + token = Math.floor(token / HEX_BASE) + } + return bytes + } + + token -= ORDINARY_OFFSET + const bytes = new Uint8Array(defaultNgramLength) + for (let i = bytes.length - 1; i >= 0; i -= 1) { + bytes[i] = ORDINARY_ALPHABET.charCodeAt(token % ORDINARY_BASE) + token = Math.floor(token / ORDINARY_BASE) + } + return bytes +} + /** * @param {string} text source text * @param {Set} out block accumulator @@ -148,6 +214,50 @@ function extractOrdinaryCodes(text, out) { } } +/** + * @param {Uint8Array} bytes UTF-8 source text + * @param {Set} out block accumulator + * @returns {void} + */ +function extractOrdinaryByteCodes(bytes, out) { + if (bytes.length < defaultNgramLength) return + let code = 0 + let count = 0 + for (let i = 0; i < bytes.length; i += 1) { + const byte = bytes[i] + let symbol = byte < ordinarySymbols.length ? ordinarySymbols[byte] : -1 + let boundaryAfter = false + if (symbol < 0 && byte === 0xc4 && bytes[i + 1] === 0xb0) { + // U+0130 lowercases to `i` followed by U+0307, which ends the run. + symbol = ordinarySymbols[105] + boundaryAfter = true + i += 1 + } else if ( + symbol < 0 && + byte === 0xe2 && + bytes[i + 1] === 0x84 && + bytes[i + 2] === 0xaa + ) { + // U+212A KELVIN SIGN lowercases to ASCII `k`. + symbol = ordinarySymbols[107] + i += 2 + } + if (symbol < 0) { + code = 0 + count = 0 + continue + } + if (count === defaultNgramLength) code %= ORDINARY_ROLL + else count += 1 + code = code * ORDINARY_BASE + symbol + if (count === defaultNgramLength) out.add(ORDINARY_OFFSET + code) + if (boundaryAfter) { + code = 0 + count = 0 + } + } +} + /** * Extract both ordinary shape lengths and hexadecimal shapes in one traversal. * The scalar rolling codes avoid allocating ring buffers and token strings. @@ -250,6 +360,105 @@ function extractStructuralCodes(text, out) { } } +/** + * @param {Uint8Array} bytes UTF-8 source text + * @param {Set} out block accumulator + * @returns {void} + */ +function extractStructuralByteCodes(bytes, out) { + if (bytes.length < SHORT_SHAPE_LENGTH) return + let longCode = 0 + let longCount = 0 + let longInformative = 0 + let shortCode = 0 + let shortCount = 0 + let shortInformative = 0 + let hexCode = 0 + let hexCount = 0 + let hexNonUniform = 0 + const uniformHexToken = HEX_OFFSET + uniformHexCode + let uniformHexAdded = out.has(uniformHexToken) + + for (let i = 0; i < bytes.length; i += 1) { + const codeUnit = bytes[i] + const digit = codeUnit >= 48 && codeUnit <= 57 + const lower = codeUnit >= 97 && codeUnit <= 122 + const upper = codeUnit >= 65 && codeUnit <= 90 + const letter = lower || upper + const punctuation = !digit && !letter && codeUnit < shapeSymbols.length + ? shapeSymbols[codeUnit] + : -1 + const shapeSymbol = digit + ? shapeDigit + : letter ? shapeLetter : punctuation + + if (shapeSymbol < 0) { + longCode = 0 + longCount = 0 + longInformative = 0 + shortCode = 0 + shortCount = 0 + shortInformative = 0 + } else { + if (longCount === SHAPE_LENGTH) { + if (Math.floor(longCode / SHAPE_ROLL) !== shapeLetter) { + longInformative -= 1 + } + longCode %= SHAPE_ROLL + } else { + longCount += 1 + } + longCode = longCode * SHAPE_BASE + shapeSymbol + if (shapeSymbol !== shapeLetter) longInformative += 1 + if (longCount === SHAPE_LENGTH && longInformative > 0) out.add(longCode) + + if (shortCount === SHORT_SHAPE_LENGTH) { + if (Math.floor(shortCode / SHORT_SHAPE_ROLL) !== shapeLetter) { + shortInformative -= 1 + } + shortCode %= SHORT_SHAPE_ROLL + } else { + shortCount += 1 + } + shortCode = shortCode * SHAPE_BASE + shapeSymbol + if (shapeSymbol !== shapeLetter) shortInformative += 1 + if (shortCount === SHORT_SHAPE_LENGTH && shortInformative > 0) { + out.add(shortCode * SHORT_SHAPE_PADDING) + } + } + + const hexLetter = lower && codeUnit <= 102 || upper && codeUnit <= 70 + const hexPunctuation = !digit && !letter && codeUnit < hexSymbols.length + ? hexSymbols[codeUnit] + : -1 + const hexSymbol = digit || hexLetter ? hexUniform : hexPunctuation + if (hexSymbol < 0 || bytes.length < HEX_LENGTH) { + hexCode = 0 + hexCount = 0 + hexNonUniform = 0 + continue + } + if (hexCount === HEX_LENGTH) { + if (Math.floor(hexCode / HEX_ROLL) !== hexUniform) hexNonUniform -= 1 + hexCode %= HEX_ROLL + } else { + hexCount += 1 + } + hexCode = hexCode * HEX_BASE + hexSymbol + if (hexSymbol !== hexUniform) hexNonUniform += 1 + if (hexCount === HEX_LENGTH) { + if (hexNonUniform === 0) { + if (!uniformHexAdded) { + out.add(uniformHexToken) + uniformHexAdded = true + } + } else { + out.add(HEX_OFFSET + hexCode) + } + } + } +} + /** * @param {number} code encoded digits * @param {number} length output length diff --git a/test/indexTokens.test.js b/test/indexTokens.test.js index b92b9b7..7d15769 100644 --- a/test/indexTokens.test.js +++ b/test/indexTokens.test.js @@ -7,7 +7,7 @@ import { defaultShapeNgramLength, defaultShortShapeNgramLength, } from '../src/constants.js' -import { decodeIndexToken, extractIndexTokenCodes, supportsNumericIndexTokens } from '../src/indexTokens.js' +import { decodeIndexToken, decodeIndexTokenBytes, extractIndexTokenBytes, extractIndexTokenCodes, supportsNumericIndexTokens } from '../src/indexTokens.js' import { extractNgrams } from '../src/ngrams.js' import { extractStructuralNgrams } from '../src/shape.js' @@ -68,6 +68,20 @@ describe('numeric index tokens', () => { expect(numeric.tokens).toEqual(strings) }) + it('extracts identical tokens directly from UTF-8 bytes', () => { + const encoder = new TextEncoder() + const texts = [ + '{"role":"USER","content":"Hello world"}', + 'emoji 🚀 and café are boundaries', + 'abCDİefgh ijklKmnop', + 'ΩЖ中 2026/07/31 test@example.com deadbeef', + ] + for (const text of texts) { + expect(extractIndexTokenBytes(encoder.encode(text))) + .toEqual(extractIndexTokenCodes(text)) + } + }) + it('preserves global string token ordering', () => { const texts = [ 'zebra alpha 90909 qqqqq m4ng0es 00abc', @@ -106,6 +120,17 @@ describe('numeric index tokens', () => { .toEqual(Array.from(strings).sort()) }) + it('decodes numeric tokens directly to their UTF-8 representation', () => { + const encoder = new TextEncoder() + const { codes } = extractNumericTokens([ + '{"role":"user"}', + '2026/07/31 test@example.com deadbeef', + ]) + for (const code of codes) { + expect(decodeIndexTokenBytes(code)).toEqual(encoder.encode(decodeIndexToken(code))) + } + }) + it('rejects invalid numeric tokens', () => { expect(() => decodeIndexToken(-1)).toThrow('Invalid numeric index token') expect(() => decodeIndexToken(Number.MAX_SAFE_INTEGER)).toThrow('Invalid numeric index token')