diff --git a/package.json b/package.json index 9e76bcf..b9c0815 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.1" + "hyparquet-writer": "0.16.3" }, "devDependencies": { "@eslint/js": "10.0.1", diff --git a/src/createIndex.js b/src/createIndex.js index 540bb46..c6e3547 100644 --- a/src/createIndex.js +++ b/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 | undefined>} */ const asciiBuckets = [] /** @type {Map>} */ const unicodeBuckets = new Map() /** @type {{ascii: boolean, first: number, id: number, second: number}[]} */ const activePrefixes = [] + /** @type {Array | undefined>} */ + const numericBuckets = [] + /** @type {number[]} */ + const activeNumericPrefixes = [] let blockId = 0 let rowsInBlock = 0 let postingRows = 0 /** @type {Set} */ let blockNgrams = new Set() + /** @type {Set} */ + 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} 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: * ordinary character n-grams plus shape (character-class skeleton) n-grams, diff --git a/src/indexTokens.js b/src/indexTokens.js new file mode 100644 index 0000000..d1ca03b --- /dev/null +++ b/src/indexTokens.js @@ -0,0 +1,280 @@ +/** + * Compact numeric tokens used only while building an index with the default + * tokenizer. The numeric order exactly matches JavaScript's lexicographic order + * for the corresponding string tokens, so sorting codes produces the same + * Parquet ngram column as sorting strings. + * + * Ordinary 5-grams use a 40-character alphabet and fit in 27 bits. Structural + * shapes use base 9, reserving zero as an end-of-token marker so four-character + * and eight-character shapes share one correctly ordered namespace. Hex shapes + * have a fixed length and use their seven-character alphabet directly. + */ +import { defaultHexShapeNgramLength, defaultNgramChars, defaultNgramLength, defaultShapeChars, defaultShapeNgramLength, defaultShortShapeNgramLength } from './constants.js' +import { HEX_SHAPE_PREFIX, SHAPE_PREFIX } from './shape.js' + +const ORDINARY_ALPHABET = '"0123456789:abcdefghijklmnopqrstuvwxyz{}' +const ORDINARY_BASE = ORDINARY_ALPHABET.length +const ORDINARY_ROLL = ORDINARY_BASE ** (defaultNgramLength - 1) +const ORDINARY_SPACE = ORDINARY_BASE ** defaultNgramLength + +// ASCII lexical order of default structural symbols. Zero is reserved for the +// end marker, so real shape symbols are encoded as their index plus one. +const SHAPE_ALPHABET = Array.from(new Set(defaultShapeChars + 'DL')).sort().join('') +const SHAPE_BASE = SHAPE_ALPHABET.length + 1 +const SHAPE_LENGTH = defaultShapeNgramLength +const SHORT_SHAPE_LENGTH = defaultShortShapeNgramLength +const SHAPE_ROLL = SHAPE_BASE ** (SHAPE_LENGTH - 1) +const SHORT_SHAPE_ROLL = SHAPE_BASE ** (SHORT_SHAPE_LENGTH - 1) +const SHORT_SHAPE_PADDING = SHAPE_BASE ** (SHAPE_LENGTH - SHORT_SHAPE_LENGTH) +const SHAPE_SPACE = SHAPE_BASE ** SHAPE_LENGTH + +const HEX_ALPHABET = Array.from(new Set(defaultShapeChars + 'H')).sort().join('') +const HEX_BASE = HEX_ALPHABET.length +const HEX_LENGTH = defaultHexShapeNgramLength +const HEX_ROLL = HEX_BASE ** (HEX_LENGTH - 1) +const HEX_SPACE = HEX_BASE ** HEX_LENGTH + +const HEX_OFFSET = SHAPE_SPACE +const ORDINARY_OFFSET = HEX_OFFSET + HEX_SPACE +const TOKEN_SPACE = ORDINARY_OFFSET + ORDINARY_SPACE + +const ordinarySymbols = new Int8Array(128).fill(-1) +for (let i = 0; i < ORDINARY_ALPHABET.length; i += 1) { + ordinarySymbols[ORDINARY_ALPHABET.charCodeAt(i)] = i +} + +const shapeSymbols = new Int8Array(128).fill(-1) +for (let i = 0; i < SHAPE_ALPHABET.length; i += 1) { + shapeSymbols[SHAPE_ALPHABET.charCodeAt(i)] = i + 1 +} +const shapeDigit = shapeSymbols[68] // D +const shapeLetter = shapeSymbols[76] // L + +const hexSymbols = new Int8Array(128).fill(-1) +for (let i = 0; i < HEX_ALPHABET.length; i += 1) { + hexSymbols[HEX_ALPHABET.charCodeAt(i)] = i +} +const hexUniform = hexSymbols[72] // H +let uniformHexCode = 0 +for (let i = 0; i < HEX_LENGTH; i += 1) { + uniformHexCode = uniformHexCode * HEX_BASE + hexUniform +} + +/** + * Return whether the build can use compact numeric tokens. Non-default options + * retain the general string implementation, including custom Unicode chars. + * + * @param {number} n n-gram length + * @param {string} chars extra characters retained inside n-gram runs + * @returns {boolean} + */ +export function supportsNumericIndexTokens(n, chars) { + return n === defaultNgramLength && chars === defaultNgramChars +} + +/** + * Add all ordinary and structural default index tokens from a string to `out`. + * + * @param {string} text source text + * @param {Set} [out] block accumulator + * @returns {Set} + */ +export function extractIndexTokenCodes(text, out = new Set()) { + if (typeof text !== 'string') return out + extractOrdinaryCodes(text, out) + extractStructuralCodes(text, out) + return out +} + +/** + * Decode a numeric build token into the exact string stored in the index. + * + * @param {number} token numeric index token + * @returns {string} + */ +export function decodeIndexToken(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) + let shape = SHAPE_PREFIX + for (const digit of digits) { + if (digit === 0) break + shape += SHAPE_ALPHABET[digit - 1] + } + return shape + } + if (token < ORDINARY_OFFSET) { + return HEX_SHAPE_PREFIX + decodeAlphabet( + token - HEX_OFFSET, + HEX_LENGTH, + HEX_BASE, + HEX_ALPHABET + ) + } + return decodeAlphabet( + token - ORDINARY_OFFSET, + defaultNgramLength, + ORDINARY_BASE, + ORDINARY_ALPHABET + ) +} + +/** + * @param {string} text source text + * @param {Set} out block accumulator + * @returns {void} + */ +function extractOrdinaryCodes(text, out) { + if (text.length < defaultNgramLength) return + const lower = text.toLowerCase() + let code = 0 + let count = 0 + for (let i = 0; i < lower.length; i += 1) { + const codeUnit = lower.charCodeAt(i) + const symbol = codeUnit < ordinarySymbols.length + ? ordinarySymbols[codeUnit] + : -1 + 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) + } +} + +/** + * Extract both ordinary shape lengths and hexadecimal shapes in one traversal. + * The scalar rolling codes avoid allocating ring buffers and token strings. + * + * @param {string} text source text + * @param {Set} out block accumulator + * @returns {void} + */ +function extractStructuralCodes(text, out) { + if (text.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 < text.length; i += 1) { + const codeUnit = text.charCodeAt(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 || text.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 + * @param {number} base numeric base + * @returns {number[]} + */ +function decodeDigits(code, length, base) { + const digits = new Array(length) + for (let i = length - 1; i >= 0; i -= 1) { + digits[i] = code % base + code = Math.floor(code / base) + } + return digits +} + +/** + * @param {number} code encoded digits + * @param {number} length output length + * @param {number} base numeric base + * @param {string} alphabet digit alphabet + * @returns {string} + */ +function decodeAlphabet(code, length, base, alphabet) { + const digits = decodeDigits(code, length, base) + let token = '' + for (const digit of digits) token += alphabet[digit] + return token +} diff --git a/test/indexTokens.test.js b/test/indexTokens.test.js new file mode 100644 index 0000000..b92b9b7 --- /dev/null +++ b/test/indexTokens.test.js @@ -0,0 +1,113 @@ +import { describe, expect, it } from 'vitest' +import { + defaultHexShapeNgramLength, + defaultNgramChars, + defaultNgramLength, + defaultShapeChars, + defaultShapeNgramLength, + defaultShortShapeNgramLength, +} from '../src/constants.js' +import { decodeIndexToken, extractIndexTokenCodes, supportsNumericIndexTokens } from '../src/indexTokens.js' +import { extractNgrams } from '../src/ngrams.js' +import { extractStructuralNgrams } from '../src/shape.js' + +/** + * Extract the string tokens used by the general index-build path. + * + * @param {string[]} texts source strings in one logical block + * @returns {Set} + */ +function extractStringTokens(texts) { + const tokens = new Set() + for (const text of texts) { + for (const token of extractNgrams(text, defaultNgramLength, defaultNgramChars)) { + tokens.add(token) + } + extractStructuralNgrams( + text, + [defaultShapeNgramLength, defaultShortShapeNgramLength], + defaultHexShapeNgramLength, + defaultShapeChars, + tokens + ) + } + return tokens +} + +/** + * Decode the numeric tokens used by the optimized index-build path. + * + * @param {string[]} texts source strings in one logical block + * @returns {{codes: Set, tokens: Set}} + */ +function extractNumericTokens(texts) { + const codes = new Set() + for (const text of texts) extractIndexTokenCodes(text, codes) + return { codes, tokens: new Set(Array.from(codes, decodeIndexToken)) } +} + +describe('numeric index tokens', () => { + it('supports exactly the default ordinary tokenizer configuration', () => { + expect(supportsNumericIndexTokens(defaultNgramLength, defaultNgramChars)).toBe(true) + expect(supportsNumericIndexTokens(defaultNgramLength + 1, defaultNgramChars)).toBe(false) + expect(supportsNumericIndexTokens(defaultNgramLength, '')).toBe(false) + }) + + it('matches ordinary, shape, short-shape, and hex string tokens', () => { + const texts = [ + '{"role":"USER","content":"Hello world"}', + 'date=2026/07/31 time=12:34:56 phone=415-555-0123', + 'email@example.com 192.168.10.42 abc-1234 abc-12345678', + '550e8400-e29b-41d4-a716-446655440000 deadbeef cafebabe', + 'Unicode Ω and café are token boundaries', + ] + const strings = extractStringTokens(texts) + const numeric = extractNumericTokens(texts) + + expect(numeric.codes.size).toBe(strings.size) + expect(numeric.tokens).toEqual(strings) + }) + + it('preserves global string token ordering', () => { + const texts = [ + 'zebra alpha 90909 qqqqq m4ng0es 00abc', + '1-234 12-345678 abcdef12-3456-7890-abcd-ef1234567890', + '{"a":12345,"z":"value"} test@example.com 2026/7/1', + ] + const strings = extractStringTokens(texts) + const { codes } = extractNumericTokens(texts) + const decodedOrder = Array.from(codes) + .sort((a, b) => a - b) + .map(decodeIndexToken) + + expect(decodedOrder).toEqual(Array.from(strings).sort()) + }) + + it('matches the string implementation across deterministic mixed input', () => { + const alphabet = 'aAfFgG09@.-_/:"{}, []() Ωé\n' + let state = 0x12345678 + const texts = [] + for (let row = 0; row < 250; row += 1) { + state = Math.imul(state, 1664525) + 1013904223 | 0 + const length = 10 + (state >>> 0) % 150 + let text = '' + for (let i = 0; i < length; i += 1) { + state = Math.imul(state, 1664525) + 1013904223 | 0 + text += alphabet[(state >>> 0) % alphabet.length] + } + texts.push(text) + } + + const strings = extractStringTokens(texts) + const numeric = extractNumericTokens(texts) + expect(numeric.codes.size).toBe(strings.size) + expect(numeric.tokens).toEqual(strings) + expect(Array.from(numeric.codes).sort((a, b) => a - b).map(decodeIndexToken)) + .toEqual(Array.from(strings).sort()) + }) + + it('rejects invalid numeric tokens', () => { + expect(() => decodeIndexToken(-1)).toThrow('Invalid numeric index token') + expect(() => decodeIndexToken(Number.MAX_SAFE_INTEGER)).toThrow('Invalid numeric index token') + }) +})