From d8353b6b32b25ea1db605f7a62e3c828f00439f5 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Wed, 1 Jul 2026 02:18:31 -0700 Subject: [PATCH] Add shape (character-class skeleton) index layer Accelerates literal-free pattern regexes (phones, SSNs, IPs, dates) that have no fixed n-gram to prune on. Text is indexed under a second, coarser tokenizer: letters collapse to L, digits to D, structural punctuation is kept, everything else is a boundary. Only shapes containing a non-letter are emitted, so a short shape stays selective (a phone DDD-DDDD is rare even though every character class is common). Queries project the regex through the same coarse alphabet and reuse the existing extractRegexLiterals, so there is no bespoke analyzer and no hand-coded assumptions about what a query looks like: /\d{3}-\d{2}-\d{4}/ becomes D{3}-D{2}-D{4} and the trusted extractor yields DDD-DD-DDDD. The index stays a superset candidate filter and the per-row scan still delivers exact grep, so shape tokens (mandatory, hence never a false negative) only change which blocks are fetched, never the results. On 25k WildChat rows this takes SSN/date/time/phone regexes from full scans to 5-26% of blocks for ~2% more index. The layer is always on and not configurable; its settings live in kv metadata (format, not config) so queries tokenize identically and shape-less indexes are detected. --- src/constants.js | 13 +++ src/createIndex.js | 31 ++++-- src/queryIndex.js | 45 +++++++-- src/regex.js | 203 ++++++++++++++++++++++++++++++++++++++- src/shape.js | 78 +++++++++++++++ src/types.d.ts | 2 + test/createIndex.test.js | 12 ++- test/queryIndex.test.js | 15 +++ test/regex.test.js | 11 +++ test/shape.test.js | 200 ++++++++++++++++++++++++++++++++++++++ 10 files changed, 586 insertions(+), 24 deletions(-) create mode 100644 src/shape.js create mode 100644 test/shape.test.js diff --git a/src/constants.js b/src/constants.js index 32f314c..1157631 100644 --- a/src/constants.js +++ b/src/constants.js @@ -29,3 +29,16 @@ export const defaultNgramLength = 5 // identically; an index with no such key (built before this) is read as the // empty set, i.e. plain alphanumeric, and keeps working unchanged. export const defaultNgramChars = '"{}:' + +// Length of shape (character-class skeleton) n-grams. Shape n-grams collapse +// letters to L and digits to D over a tiny alphabet, so a short window stays +// selective for structured patterns (a phone shape DDD-DDDD is rare even though +// every individual character class is common). 4 captures the digit-group +// structure of phones, SSNs, IPs and dates. +export const defaultShapeNgramLength = 4 + +// Structural punctuation kept verbatim inside shape runs (everything else that +// is not a letter or digit is a shape boundary). These are the separators that +// give numeric patterns their selectivity: dots, dashes, slashes, colons, +// underscores and the email '@'. +export const defaultShapeChars = '@.-_/:' diff --git a/src/createIndex.js b/src/createIndex.js index 3e4aec0..1d3eeeb 100644 --- a/src/createIndex.js +++ b/src/createIndex.js @@ -1,7 +1,8 @@ import { parquetMetadataAsync, parquetReadObjects } from 'hyparquet' import { ParquetWriter, schemaFromColumnData } from 'hyparquet-writer' -import { defaultBlockSize, defaultIndexRowGroupSize, defaultNgramChars, defaultNgramLength, hypGrepVersion } from './constants.js' +import { defaultBlockSize, defaultIndexRowGroupSize, defaultNgramChars, defaultNgramLength, defaultShapeChars, defaultShapeNgramLength, hypGrepVersion } from './constants.js' import { extractNgrams } from './ngrams.js' +import { extractShapeNgrams } from './shape.js' import { assertNonNegativeSafeInteger, assertPositiveSafeInteger, getTextColumnsFromSchema } from './utils.js' /** @@ -32,6 +33,12 @@ export async function createIndex({ } assertNonNegativeSafeInteger(sourceFile.byteLength, 'sourceFile.byteLength') + // The shape layer is always on and not configurable; it is a second, coarser + // tokenization of the same text. Its settings are recorded in kv metadata so + // queries tokenize identically and so an index built without it is detected. + const shapeNgramLength = defaultShapeNgramLength + const shapeChars = defaultShapeChars + const metadata = sourceMetadata ?? await parquetMetadataAsync(sourceFile) const numRows = Number(metadata.num_rows) assertNonNegativeSafeInteger(numRows, 'source row count') @@ -103,7 +110,7 @@ export async function createIndex({ columns: textColumns, }) for (const row of rows) { - collectRowNgrams(row, textColumns, ngramLength, ngramChars, blockNgrams) + collectRowNgrams(row, textColumns, ngramLength, ngramChars, shapeNgramLength, shapeChars, blockNgrams) if (++rowsInBlock === blockSize) flushBlock() } groupStart += groupRows @@ -118,6 +125,10 @@ export async function createIndex({ // missing this key (built before structural tokenization) is read as '', // i.e. plain alphanumeric, and keeps working unchanged. { key: 'hypgrep.ngram_chars', value: ngramChars }, + // Shape (character-class skeleton) layer. Absence of these keys means the + // index has no shape tokens, so queryIndex must not require any. + { key: 'hypgrep.shape_ngram_length', value: String(shapeNgramLength) }, + { key: 'hypgrep.shape_chars', value: shapeChars }, { key: 'hypgrep.text_columns', value: textColumns.join(',') }, { key: 'hypgrep.source_rows', value: String(numRows) }, // Can save network requests on the source file @@ -183,21 +194,27 @@ export async function createIndex({ } /** - * Collect the distinct n-grams of one row into a block's accumulator set. + * 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, + * which share the column but are namespaced by SHAPE_PREFIX. * * @param {Record} row * @param {string[]} textColumns * @param {number} n * @param {string} chars extra characters kept inside n-gram runs + * @param {number} shapeN shape n-gram length + * @param {string} shapeChars structural punctuation kept verbatim in shapes * @param {Set} ngrams block accumulator to add into */ -function collectRowNgrams(row, textColumns, n, chars, ngrams) { +function collectRowNgrams(row, textColumns, n, chars, shapeN, shapeChars, ngrams) { if (!row) return for (const columnName of textColumns) { const value = row[columnName] - if (typeof value !== 'string' || value.length < n) continue - for (const g of extractNgrams(value, n, chars)) { - ngrams.add(g) + if (typeof value !== 'string') continue + if (value.length >= n) { + for (const g of extractNgrams(value, n, chars)) ngrams.add(g) } + // Shape grams stream straight into the block accumulator (no throwaway set). + extractShapeNgrams(value, shapeN, shapeChars, ngrams) } } diff --git a/src/queryIndex.js b/src/queryIndex.js index 62216e9..4eec1c6 100644 --- a/src/queryIndex.js +++ b/src/queryIndex.js @@ -1,7 +1,7 @@ import { parquetMetadataAsync, parquetQuery } from 'hyparquet' import { hypGrepVersion } from './constants.js' import { literalsToNgrams, queryNgrams } from './ngrams.js' -import { extractRegexLiterals } from './regex.js' +import { extractRegexLiterals, extractRegexShapes } from './regex.js' import { assertNonNegativeSafeInteger, assertPositiveSafeInteger } from './utils.js' /** @@ -32,19 +32,31 @@ export async function queryIndex({ query, indexFile, indexMetadata }) { // Read index kv metadata indexMetadata ??= await parquetMetadataAsync(indexFile) const kvMetadata = indexMetadata.key_value_metadata || [] - const { blockSize, ngramLength, ngramChars, textColumns, sourceByteLength, sourceRows } = parseKvMetadata(kvMetadata) + const { blockSize, ngramLength, ngramChars, shapeNgramLength, shapeChars, textColumns, sourceByteLength, sourceRows } = parseKvMetadata(kvMetadata) // A "branch" is a conjunction of n-grams that ALL must appear in a block. // A query matches a block if ANY branch is fully satisfied (DNF). // // - String query: one branch, the n-grams of the string itself. - // - RegExp query: one branch per top-level alternation arm. - // The query MUST be tokenized with the same kept-character set the index was - // built with, or its n-grams won't line up with the postings. + // - RegExp query: one branch per top-level alternation arm, plus any mandatory + // shape n-grams (character-class skeletons) ANDed in. Shape tokens let a + // literal-free pattern (a phone number, SSN, IP) prune by its structure. + // Gated on shapeNgramLength so an index with no shape layer is never asked + // for shape tokens (which would be a false negative). /** @type {string[][]} */ - const branches = query instanceof RegExp - ? extractRegexLiterals(query).map(lits => literalsToNgrams(lits, ngramLength, ngramChars)) - : [queryNgrams(query, ngramLength, ngramChars)] + let branches + if (query instanceof RegExp) { + const shapeGrams = shapeNgramLength ? extractRegexShapes(query, shapeNgramLength, shapeChars) : [] + branches = extractRegexLiterals(query).map(lits => { + const grams = literalsToNgrams(lits, ngramLength, ngramChars) + // Shape grams are near-universal (any 4-digit run, any hyphenated word), + // so their posting lists cost bytes but prune little. Only spend them on a + // branch with no literal n-gram to prune on; otherwise the literal wins. + return grams.length ? grams : grams.concat(shapeGrams) + }) + } else { + branches = [queryNgrams(query, ngramLength, ngramChars)] + } // If any branch is empty, that branch matches anything — falling back to a // full scan is correct and bounded. @@ -124,6 +136,11 @@ export function parseKvMetadata(kvMetadata) { // Extra characters kept inside n-gram runs. Absent in pre-structural indexes, // which were tokenized as plain alphanumeric, so '' reproduces that exactly. let ngramChars = '' + // Shape layer settings. shapeNgramLength stays undefined for indexes built + // without a shape layer, which signals queryIndex not to require shape tokens. + /** @type {number | undefined} */ + let shapeNgramLength + let shapeChars = '' /** @type {string[]} */ let textColumns = [] /** @type {number | undefined} */ @@ -141,6 +158,16 @@ export function parseKvMetadata(kvMetadata) { if (key === 'hypgrep.ngram_chars' && typeof value === 'string') { ngramChars = value } + if (key === 'hypgrep.shape_ngram_length') { + // The shape layer is optional: a malformed value degrades to "no shape + // layer" (like a missing key) instead of failing the whole index — the + // text n-gram layer is still intact and queryable. + const n = Number(value) + if (Number.isSafeInteger(n) && n > 0) shapeNgramLength = n + } + if (key === 'hypgrep.shape_chars' && typeof value === 'string') { + shapeChars = value + } if (key === 'hypgrep.version') { version = Number(value) if (version !== hypGrepVersion) { @@ -180,5 +207,5 @@ export function parseKvMetadata(kvMetadata) { assertNonNegativeSafeInteger(sourceRows, 'hypgrep.source_rows') assertNonNegativeSafeInteger(sourceByteLength, 'hypgrep.source_bytelength') - return { blockSize, ngramLength, ngramChars, textColumns, sourceByteLength, sourceRows } + return { blockSize, ngramLength, ngramChars, shapeNgramLength, shapeChars, textColumns, sourceByteLength, sourceRows } } diff --git a/src/regex.js b/src/regex.js index b12ba8f..7f87e86 100644 --- a/src/regex.js +++ b/src/regex.js @@ -1,3 +1,5 @@ +import { SHAPE_PREFIX, shapeSymbol } from './shape.js' + /** * Extract a disjunctive-normal-form set of mandatory literal substrings from a * regex. The return value is an array of "branches"; each branch is the array @@ -35,12 +37,201 @@ * @returns {string[][]} */ export function extractRegexLiterals(regex) { + // v-mode class syntax (nested classes, set operations) desyncs skipClass, + // which would turn class members into phantom mandatory literals — a false + // negative. One unprunable branch means the caller falls back to a full scan. + if (regex.flags.includes('v')) return [[]] const lower = regex.flags.includes('i') return expandRegion(regex.source, 0, regex.source.length, lower) } const MAX_DNF = 32 +/** + * Extract mandatory shape n-grams from a regex, for the shape (character-class + * skeleton) index layer. Lets patterns with no fixed literal — phone numbers, + * SSNs, IPs, dates — still prune blocks by their digit/punctuation structure. + * + * Rather than reason about the regex directly, we project it into the coarse + * shape alphabet (letters → L, digits → D, punctuation kept) and hand the + * result to {@link extractRegexLiterals}. The existing extractor already knows + * how to find the mandatory literals through quantifiers, classes and groups, + * so the only shape-specific code is the mechanical projection; there are no + * hand-coded assumptions about what the query "looks like". + * + * Shape tokens are only required when they are mandatory across the WHOLE + * pattern, i.e. the projection yields a single branch. A projected alternation + * (more than one branch) means no shape is common to every match, so we add + * none and let the literal branches prune on their own — conservative, never a + * false negative. + * + * @param {RegExp} regex + * @param {number} shapeN shape n-gram length + * @param {string} shapeChars structural punctuation kept verbatim in shapes + * @returns {string[]} prefixed shape n-grams (all mandatory; empty if none) + */ +export function extractRegexShapes(regex, shapeN, shapeChars) { + // v-mode class syntax desyncs the scanner (see extractRegexLiterals), and + // i+u case folding lets an ASCII letter atom match non-ASCII characters + // (U+017F ſ, U+212A K) that the index maps to run boundaries. Either would + // emit shape grams that aren't truly mandatory — a false negative — so + // require no shapes at all. + const { flags } = regex + if (flags.includes('v') || flags.includes('i') && flags.includes('u')) return [] + /** @type {string[][]} */ + let branches + try { + // No 'i' flag: shape symbols L/D must not be lowercased. + branches = extractRegexLiterals(new RegExp(projectToShape(regex.source, shapeChars))) + } catch { + return [] + } + if (branches.length !== 1) return [] + /** @type {Set} */ + const out = new Set() + for (const lit of branches[0]) { + for (let i = 0; i + shapeN <= lit.length; i += 1) { + const gram = lit.slice(i, i + shapeN) + if (!/^L+$/.test(gram)) out.add(SHAPE_PREFIX + gram) + } + } + return [...out] +} + +/** + * Rewrite a regex source into the coarse shape alphabet, atom by atom, leaving + * quantifiers, anchors, alternation and group structure for the literal + * extractor to interpret. Each consuming atom becomes a single shape token: + * letters → L, digits → D, kept punctuation → that (escaped) literal, and + * anything that isn't a single shape symbol (a group, `.`, `\w`, a mixed class) + * becomes `.` so the extractor treats it as a run boundary. + * + * @param {string} src regex source + * @param {string} shapeChars structural punctuation kept verbatim in shapes + * @returns {string} a projected regex source + */ +function projectToShape(src, shapeChars) { + let out = '' + let i = 0 + while (i < src.length) { + const c = src[i] + // Pass anchors and alternation through unchanged. + if (c === '^' || c === '$' || c === '|') { out += c; i += 1; continue } + + /** @type {string | null} */ + let sym + let atomEnd + if (c === '\\') { + const next = src[i + 1] + if (next === undefined) { i += 1; continue } + if (next === 'd') { sym = 'D'; atomEnd = i + 2 } + else if (/[a-zA-Z0-9]/.test(next)) { sym = null; atomEnd = skipSpecialEscape(src, i) } // \w \s \D backref \p{} + else { sym = shapeSymbol(next, shapeChars); atomEnd = i + 2 } // \. \- \/ escaped literal + } else if (c === '[') { + atomEnd = skipClass(src, i) + sym = classShapeSymbol(src, i, atomEnd, shapeChars) + } else if (c === '(') { + atomEnd = skipGroup(src, i) + const group = parseGroupPrefix(src, i) + if (group.kind === 'consuming') { + // Project the group body in place and keep the (...) structure so the + // literal extractor folds it — grouped patterns like (\d{3})-(\d{4}) + // still prune by shape instead of collapsing to a boundary. + const body = projectToShape(src.slice(group.innerStart, atomEnd - 1), shapeChars) + const q = peekQuantifier(src, atomEnd) + out += src.slice(i, group.innerStart) + body + ')' + src.slice(atomEnd, q.end) + i = q.end + continue + } + sym = null // assertion / unknown (?...) → boundary + } else if (c === '.') { + atomEnd = i + 1 + sym = null + } else { + atomEnd = i + 1 + sym = shapeSymbol(c, shapeChars) + } + + // Copy any quantifier on the atom verbatim so the extractor expands it. + const q = peekQuantifier(src, atomEnd) + out += shapeAtomSource(sym) + src.slice(atomEnd, q.end) + i = q.end + } + return out +} + +/** + * Render a shape symbol as regex source: `.` for a boundary (null), L/D as-is, + * and a kept punctuation char as an escaped literal. + * + * @param {string | null} sym + * @returns {string} + */ +function shapeAtomSource(sym) { + if (sym === null) return '.' + if (sym === 'L' || sym === 'D') return sym + return /[.*+?^${}()|[\]\\/]/.test(sym) ? '\\' + sym : sym +} + +/** + * Resolve a `[...]` character class to a single shape symbol, or null if its + * members don't all share one (so it can't extend a mandatory shape run). + * Members may be single characters, `\d`, an escaped punctuation char, or a + * plain range like `0-5` or `2-9` whose endpoints sit in the same band. + * + * @param {string} src + * @param {number} start index of '[' + * @param {number} classEnd index after ']' + * @param {string} shapeChars + * @returns {string | null} + */ +function classShapeSymbol(src, start, classEnd, shapeChars) { + let i = start + 1 + const lastInside = classEnd - 1 // position of ']' + if (src[i] === '^') return null + /** @type {string | null} */ + let sym = null + while (i < lastInside) { + /** @type {string | null} */ + let s + if (src[i] === '\\') { + const next = src[i + 1] + if (next === undefined) return null + if (next === 'd') s = 'D' + else if (/[a-zA-Z0-9]/.test(next)) return null // \w \s \x41 etc. + else s = shapeSymbol(next, shapeChars) + i += 2 + } else if (src[i + 1] === '-' && i + 2 < lastInside) { + s = rangeShapeSymbol(src[i], src[i + 2]) + i += 3 + } else { + s = shapeSymbol(src[i], shapeChars) + i += 1 + } + if (s === null) return null + if (sym === null) sym = s + else if (sym !== s) return null + } + return sym +} + +/** + * Shape symbol for a plain class range `lo-hi`, or null. A range collapses to + * one symbol only when both endpoints sit in the same band (digits, lowercase + * letters, or uppercase letters), so every character between them shares it — + * `A-z` spans punctuation and must not resolve. + * + * @param {string} lo + * @param {string} hi + * @returns {string | null} + */ +function rangeShapeSymbol(lo, hi) { + if (lo >= '0' && lo <= '9' && hi >= '0' && hi <= '9') return 'D' + if (lo >= 'a' && lo <= 'z' && hi >= 'a' && hi <= 'z') return 'L' + if (lo >= 'A' && lo <= 'Z' && hi >= 'A' && hi <= 'Z') return 'L' + return null +} + /** * Expand a (sub-)region of regex source into a DNF of literal branches. * Splits on top-level alternation first, then expands each branch. @@ -352,15 +543,21 @@ function parseEscapedLiteral(src, start, lower) { */ function skipSpecialEscape(src, start) { const next = src[start + 1] - if (next === 'x') return Math.min(start + 4, src.length) + // \x, \u and \c only consume their payload when it is well-formed. An + // invalid payload (Annex B: the escape matches the literal 'x'/'u'/'c') + // must be rescanned as the literal atoms it is, or the scan desyncs and + // later escapes are misread — producing phantom mandatory literals. + if (next === 'x') { + return /^[0-9a-fA-F]{2}$/.test(src.slice(start + 2, start + 4)) ? start + 4 : start + 2 + } if (next === 'u') { if (src[start + 2] === '{') { const close = src.indexOf('}', start + 3) return close >= 0 ? close + 1 : src.length } - return Math.min(start + 6, src.length) + return /^[0-9a-fA-F]{4}$/.test(src.slice(start + 2, start + 6)) ? start + 6 : start + 2 } - if (next === 'c') return Math.min(start + 3, src.length) + if (next === 'c') return /^[a-zA-Z]$/.test(src[start + 2] ?? '') ? start + 3 : start + 2 if ((next === 'p' || next === 'P') && src[start + 2] === '{') { const close = src.indexOf('}', start + 3) return close >= 0 ? close + 1 : src.length diff --git a/src/shape.js b/src/shape.js new file mode 100644 index 0000000..e2fd25a --- /dev/null +++ b/src/shape.js @@ -0,0 +1,78 @@ +/** + * Shape (character-class skeleton) extraction for accelerating pattern regexes + * that have no fixed literal of n-gram length — phone numbers, SSNs, IPs, dates. + * + * Every letter collapses to `L`, every digit to `D`, and a small set of + * structural punctuation characters is kept verbatim; all other characters are + * run boundaries. N-grams are then taken over this tiny alphabet. We only emit + * shapes that contain a non-letter (a digit or kept punctuation): all-letter + * shapes are as non-selective as prose itself and would only bloat the index. + * + * Shape tokens share the index's `ngram` column with ordinary character + * n-grams. They are prefixed with `SHAPE_PREFIX` (a control byte that never + * occurs in lowercased text n-grams) so the two namespaces can never collide. + * + * Index and query MUST use the same shape length and kept-character set, or the + * tokens won't line up; createIndex records both in kv metadata and queryIndex + * reads them back. An index built without a shape layer simply has no shape + * tokens, and queryIndex must not require any (it gates on the metadata). + */ + +// Control byte that namespaces shape tokens away from text n-grams. +export const SHAPE_PREFIX = String.fromCharCode(1) + +/** + * Map a single character to its shape symbol, or null if it is a run boundary. + * + * @param {string} ch single character + * @param {string} shapeChars structural punctuation kept verbatim in shapes + * @returns {string | null} + */ +export function shapeSymbol(ch, shapeChars) { + if (ch >= 'a' && ch <= 'z' || ch >= 'A' && ch <= 'Z') return 'L' + if (ch >= '0' && ch <= '9') return 'D' + if (shapeChars.includes(ch)) return ch + return null +} + +/** + * Extract the distinct, selective shape n-grams in a string into `out`. + * + * A fixed-size ring buffer holds the last `shapeN` shape symbols of the current + * run, so windows are emitted in one pass with no per-character run string and + * no per-window slice. `nonLetters` tracks how many of the buffered symbols are + * not `L`, so all-letter windows (non-selective prose) are skipped without + * building their gram. Grams go straight into `out`, avoiding a throwaway set. + * + * @param {string} text + * @param {number} shapeN shape n-gram length + * @param {string} shapeChars structural punctuation kept verbatim in shapes + * @param {Set} [out] accumulator to add grams into (default: fresh set) + * @returns {Set} + */ +export function extractShapeNgrams(text, shapeN, shapeChars, out = new Set()) { + if (typeof text !== 'string' || text.length < shapeN) return out + /** @type {string[]} ring buffer of the last shapeN shape symbols */ + const win = new Array(shapeN) + let count = 0 // symbols seen in the current run (since the last boundary) + let nonLetters = 0 // non-'L' symbols among the buffered window + for (let i = 0; i < text.length; i += 1) { + const sym = shapeSymbol(text[i], shapeChars) + if (sym === null) { count = 0; nonLetters = 0; continue } + const slot = count % shapeN + if (count >= shapeN && win[slot] !== 'L') nonLetters -= 1 + win[slot] = sym + if (sym !== 'L') nonLetters += 1 + count += 1 + // Once the window is full, emit it unless it is all-letters (prose). + if (count >= shapeN && nonLetters > 0) { + let gram = SHAPE_PREFIX + for (let k = count % shapeN, n = 0; n < shapeN; n += 1) { + gram += win[k] + k = k + 1 === shapeN ? 0 : k + 1 + } + out.add(gram) + } + } + return out +} diff --git a/src/types.d.ts b/src/types.d.ts index 5257ccb..45fabc2 100644 --- a/src/types.d.ts +++ b/src/types.d.ts @@ -64,6 +64,8 @@ export interface HypGrepMetadata { blockSize: number // number of rows per logical block ngramLength: number // n-gram size used to build the index ngramChars: string // extra characters kept inside n-gram runs ('' for pre-structural indexes) + shapeNgramLength?: number // shape n-gram size, or undefined if the index has no shape layer + shapeChars: string // structural punctuation kept verbatim in shapes ('' if no shape layer) textColumns: string[] // list of indexed text columns sourceRows: number // number of rows in the source parquet file sourceByteLength: number // byte length of the source parquet file diff --git a/test/createIndex.test.js b/test/createIndex.test.js index 0b17509..0d78712 100644 --- a/test/createIndex.test.js +++ b/test/createIndex.test.js @@ -27,19 +27,21 @@ describe('createIndex', () => { expect(existsSync(TEST_INDEX)).toBe(true) const indexBuffer = await asyncBufferFromFile(TEST_INDEX) - expect(indexBuffer.byteLength).toBe(2522) + expect(indexBuffer.byteLength).toBe(2584) const indexMetadata = await parquetMetadataAsync(indexBuffer) expect(indexMetadata.row_groups.length).toBe(7) expect(indexMetadata.num_rows).toBe(676n) - expect(indexMetadata.key_value_metadata?.length).toBe(7) + expect(indexMetadata.key_value_metadata?.length).toBe(9) const kv = indexMetadata.key_value_metadata expect(kv?.[0]).toEqual({ key: 'hypgrep.version', value: '0' }) expect(kv?.[1]).toEqual({ key: 'hypgrep.block_size', value: '200' }) expect(kv?.[2]).toEqual({ key: 'hypgrep.ngram_length', value: '5' }) expect(kv?.[3]).toEqual({ key: 'hypgrep.ngram_chars', value: '"{}:' }) - expect(kv?.[4]).toEqual({ key: 'hypgrep.text_columns', value: 'id' }) - expect(kv?.[5]).toEqual({ key: 'hypgrep.source_rows', value: '676' }) - expect(kv?.[6]).toEqual({ key: 'hypgrep.source_bytelength', value: String(sourceFile.byteLength) }) + expect(kv?.[4]).toEqual({ key: 'hypgrep.shape_ngram_length', value: '4' }) + expect(kv?.[5]).toEqual({ key: 'hypgrep.shape_chars', value: '@.-_/:' }) + expect(kv?.[6]).toEqual({ key: 'hypgrep.text_columns', value: 'id' }) + expect(kv?.[7]).toEqual({ key: 'hypgrep.source_rows', value: '676' }) + expect(kv?.[8]).toEqual({ key: 'hypgrep.source_bytelength', value: String(sourceFile.byteLength) }) }) it('should reject invalid sizing options', async () => { diff --git a/test/queryIndex.test.js b/test/queryIndex.test.js index e917e1f..ee33ade 100644 --- a/test/queryIndex.test.js +++ b/test/queryIndex.test.js @@ -75,6 +75,21 @@ describe('queryIndex', () => { expect(() => parseKvMetadata(kv)).toThrow('hypgrep.source_rows must be a non-negative safe integer') }) + it('treats a malformed shape_ngram_length as no shape layer', () => { + for (const value of ['', 'null', '0', undefined]) { + const kv = [ + { key: 'hypgrep.version', value: '0' }, + { key: 'hypgrep.block_size', value: '100' }, + { key: 'hypgrep.ngram_length', value: '5' }, + { key: 'hypgrep.shape_ngram_length', value }, + { key: 'hypgrep.text_columns', value: 'id' }, + { key: 'hypgrep.source_rows', value: '10' }, + { key: 'hypgrep.source_bytelength', value: '100' }, + ] + expect(parseKvMetadata(kv).shapeNgramLength).toBeUndefined() + } + }) + it('rejects missing index format metadata', () => { const kv = [ { key: 'hypgrep.text_columns', value: 'id' }, diff --git a/test/regex.test.js b/test/regex.test.js index c4bfc8c..c353642 100644 --- a/test/regex.test.js +++ b/test/regex.test.js @@ -6,6 +6,17 @@ describe('extractRegexLiterals', () => { expect(extractRegexLiterals(/foo/)).toEqual([['foo']]) }) + it('treats v-flag patterns as unprunable (nested class syntax)', () => { + // matches the single char 'a' or 'x' etc. — class members must not + // become mandatory literals + expect(extractRegexLiterals(/[[abc]xfooooo]/v)).toEqual([[]]) + }) + + it('rescans invalid escape payloads as the literals they match', () => { + // Annex B: invalid \u matches literal 'u', so 'grepx' is mandatory + expect(extractRegexLiterals(new RegExp('\\u12\\dgrepx'))).toEqual([['12', 'grepx']]) + }) + it('lowercases under the i flag', () => { expect(extractRegexLiterals(/Foo/i)).toEqual([['foo']]) expect(extractRegexLiterals(/Foo/)).toEqual([['Foo']]) diff --git a/test/shape.test.js b/test/shape.test.js new file mode 100644 index 0000000..1aea1ac --- /dev/null +++ b/test/shape.test.js @@ -0,0 +1,200 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { asyncBufferFromFile } from 'hyparquet' +import { fileWriter, parquetWriteFile } from 'hyparquet-writer' +import { existsSync, unlinkSync } from 'fs' +import { createIndex } from '../src/createIndex.js' +import { parquetFind } from '../src/parquetFind.js' +import { queryIndex } from '../src/queryIndex.js' +import { extractRegexShapes } from '../src/regex.js' +import { SHAPE_PREFIX, extractShapeNgrams } from '../src/shape.js' + +const SHAPE_CHARS = '@.-_/:' +const P = SHAPE_PREFIX + +describe('extractRegexShapes', () => { + it('extracts mandatory shape n-grams from an SSN pattern', () => { + const grams = extractRegexShapes(/\d{3}-\d{2}-\d{4}/, 4, SHAPE_CHARS) + // shape run "DDD-DD-DDDD" → distinct non-all-letter 4-grams + expect(new Set(grams)).toEqual(new Set([ + P + 'DDD-', P + 'DD-D', P + 'D-DD', P + '-DD-', P + '-DDD', P + 'DDDD', + ])) + }) + + it('handles literal digits and kept punctuation', () => { + // phone with literal separators and a literal area-code digit + const grams = extractRegexShapes(/1-\d{3}-\d{4}/, 4, SHAPE_CHARS) + expect(grams).toContain(P + 'D-DD') + expect(grams.every(g => g.startsWith(P))).toBe(true) + }) + + it('breaks the run on open-ended quantifiers', () => { + // \d{1,3} is mandatory for only one digit, so no 4-long numeric run survives + expect(extractRegexShapes(/\d{1,3}\.\d{1,3}/, 4, SHAPE_CHARS)).toEqual([]) + }) + + it('bails on top-level alternation (shape could differ per arm)', () => { + expect(extractRegexShapes(/\d{4}-\d{4}|abc/, 4, SHAPE_CHARS)).toEqual([]) + }) + + it('returns nothing when there is no mandatory non-letter shape', () => { + expect(extractRegexShapes(/\w+/, 4, SHAPE_CHARS)).toEqual([]) + expect(extractRegexShapes(/hello/, 4, SHAPE_CHARS)).toEqual([]) // all-letter shapes dropped + }) + + it('extracts shapes through capture groups', () => { + // grouped SSN — the common real-world spelling — prunes like the flat form + const flat = new Set(extractRegexShapes(/\d{3}-\d{2}-\d{4}/, 4, SHAPE_CHARS)) + const grouped = new Set(extractRegexShapes(/(\d{3})-(\d{2})-(\d{4})/, 4, SHAPE_CHARS)) + expect(grouped).toEqual(flat) + // named and non-capturing groups too + expect(extractRegexShapes(/(?\d{3})-\d{4}/, 4, SHAPE_CHARS)).toContain(P + 'DDD-') + expect(extractRegexShapes(/(?:\d{3})-\d{4}/, 4, SHAPE_CHARS)).toContain(P + 'DDD-') + }) + + it('resolves partial digit ranges in classes', () => { + // IP octet / NANP area-code spellings must keep their digit shape runs + expect(extractRegexShapes(/25[0-5]-\d/, 4, SHAPE_CHARS)).toContain(P + 'DDD-') + expect(extractRegexShapes(/[2-9]\d{2}-\d{4}/, 4, SHAPE_CHARS)).toContain(P + 'DDD-') + // A-z spans punctuation between the letter bands: must not resolve to L + expect(extractRegexShapes(/[A-z]1-23/, 4, SHAPE_CHARS)).toEqual([P + 'D-DD']) + }) + + it('requires no shapes for v-flag classes (scanner cannot parse them)', () => { + // matches the single char 'x', so DDDD must not become mandatory + expect(extractRegexShapes(/[[abc]x\d\d\d\d]/v, 4, SHAPE_CHARS)).toEqual([]) + }) + + it('requires no shapes under i+u case folding', () => { + // /ss-\d\d/iu matches 'ſſ-12' whose letters are shape boundaries + expect(extractRegexShapes(/ss-\d\d/iu, 4, SHAPE_CHARS)).toEqual([]) + // i alone folds within ASCII only and keeps shape pruning + expect(extractRegexShapes(/\d{3}-\d{2}-\d{4}/i, 4, SHAPE_CHARS)).toHaveLength(6) + }) + + it('rescans invalid \\x and \\u escape payloads as literals', () => { + // Annex B: invalid \u matches literal 'u', so this matches 'u12-34-56' + const grams = extractRegexShapes(new RegExp('\\u12-\\d\\d-\\d\\d'), 4, SHAPE_CHARS) + const have = extractShapeNgrams('u12-34-56', 4, SHAPE_CHARS) + for (const g of grams) expect(have.has(g)).toBe(true) + expect(grams).toContain(P + 'DD-D') + }) + + it('only matches against text that actually contains the shape', () => { + const grams = new Set(extractShapeNgrams('call 415-555-1234 today', 4, SHAPE_CHARS)) + for (const g of extractRegexShapes(/\d{3}-\d{4}/, 4, SHAPE_CHARS)) { + expect(grams.has(g)).toBe(true) + } + }) +}) + +const SRC = 'test/files/shape.source.parquet' +const IDX = 'test/files/shape.index.parquet' + +// One row per block. Three rows carry real PII; the rest are prose with scattered +// digits that must NOT shape-match an SSN / phone / IP pattern. +function writeSource() { + const rows = [] + for (let i = 0; i < 20; i += 1) { + if (i === 4) rows.push('my ssn is 123-45-6789 please keep it private') + else if (i === 9) rows.push('reach me at 415-555-0137 any time') + else if (i === 15) rows.push('the server lives at 10.0.12.255 on the lan') + else rows.push(`order ${i} shipped on day ${i * 3} with code A${i}B`) + } + parquetWriteFile({ filename: SRC, columnData: [{ name: 'text', data: rows }] }) +} + +/** + * @param {RegExp} query + * @returns {Promise} sorted matching row indices + */ +async function find(query) { + const rows = [] + for await (const row of parquetFind({ + url: SRC, + sourceFile: await asyncBufferFromFile(SRC), + indexFile: await asyncBufferFromFile(IDX), + query, + })) rows.push(row.__index__) + return rows.sort((a, b) => a - b) +} + +/** + * @param {RegExp} query + * @returns {Promise} + */ +async function candBlocks(query) { + const r = await queryIndex({ query, indexFile: await asyncBufferFromFile(IDX) }) + return r ? r.blocks.length : 0 +} + +describe('shape layer (literal-free pattern regexes)', () => { + afterEach(() => { + for (const f of [SRC, IDX]) if (existsSync(f)) unlinkSync(f) + }) + + it('prunes an SSN regex to the one matching block and returns the right row', async () => { + writeSource() + await createIndex({ sourceFile: await asyncBufferFromFile(SRC), indexFile: fileWriter(IDX), blockSize: 1 }) + + // Without the shape layer this regex has no literal and scans all 20 blocks. + expect(await candBlocks(/\d{3}-\d{2}-\d{4}/)).toBe(1) + expect(await find(/\d{3}-\d{2}-\d{4}/)).toEqual([4]) + }) + + it('prunes a phone regex and an IP regex by their shapes', async () => { + writeSource() + await createIndex({ sourceFile: await asyncBufferFromFile(SRC), indexFile: fileWriter(IDX), blockSize: 1 }) + + expect(await candBlocks(/\d{3}-\d{3}-\d{4}/)).toBeLessThan(20) + expect(await find(/\d{3}-\d{3}-\d{4}/)).toEqual([9]) + + expect(await find(/\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}/)).toEqual([15]) + }) + + it('returns exactly the brute-force result set (no false negatives)', async () => { + writeSource() + await createIndex({ sourceFile: await asyncBufferFromFile(SRC), indexFile: fileWriter(IDX), blockSize: 1 }) + + const re = /\d{3}-\d{2}-\d{4}/ + const all = await parquetReadAll() + const brute = all.map((t, i) => re.test(t) ? i : -1).filter(i => i >= 0).sort((a, b) => a - b) + expect(await find(re)).toEqual(brute) + }) + + it('stays exact for a regex with both a literal and a shape', async () => { + writeSource() + await createIndex({ sourceFile: await asyncBufferFromFile(SRC), indexFile: fileWriter(IDX), blockSize: 1 }) + // 'please' prunes to block 4 on its own; the SSN shape is not spent, and the + // combined query still returns exactly the matching row. + expect(await find(/\d{3}-\d{2}-\d{4} please/)).toEqual([4]) + }) + + it('prunes to zero blocks when the shape uses punctuation no row has', async () => { + writeSource() + await createIndex({ sourceFile: await asyncBufferFromFile(SRC), indexFile: fileWriter(IDX), blockSize: 1 }) + // slash-separated date shape DDDD/DD/DD; no row contains a '/' + expect(await candBlocks(/\d{4}\/\d{2}\/\d{2}/)).toBe(0) + expect(await find(/\d{4}\/\d{2}\/\d{2}/)).toEqual([]) + }) + + it('keeps grep exact even when a shape over-matches (subset false positives)', async () => { + writeSource() + await createIndex({ sourceFile: await asyncBufferFromFile(SRC), indexFile: fileWriter(IDX), blockSize: 1 }) + // credit-card shape's 4-grams are a subset of the SSN/phone shapes, so those + // blocks are candidates, but the per-row regex rejects them: exact result. + expect(await candBlocks(/\d{4}-\d{4}-\d{4}-\d{4}/)).toBeGreaterThan(0) + expect(await find(/\d{4}-\d{4}-\d{4}-\d{4}/)).toEqual([]) + }) +}) + +/** + * Read the source text column back as a plain array for brute-force comparison. + * + * @returns {Promise} + */ +async function parquetReadAll() { + const { parquetReadObjects } = await import('hyparquet') + const file = await asyncBufferFromFile(SRC) + const rows = await parquetReadObjects({ file, columns: ['text'] }) + return rows.map(r => r.text) +}