feat: array-unnest groupBy for aggregates + batch-embed entity extraction

- groupBy now supports { field, unnest: true }: an entity contributes once per
  distinct element of an array field (tag frequency / faceted counts). Duplicate
  elements on one entity count once; an empty/missing array joins no group. The
  incremental add/update/delete paths fan out across the unnested groups.
- extractEntities/extractConcepts batch-embed the unique candidate spans in one
  embedBatch call instead of one embed() per candidate (N sequential model calls);
  falls back to per-candidate embedding if the batch fails. No behavior change.
This commit is contained in:
David Snelling 2026-05-26 14:20:40 -07:00
parent 2591001bd0
commit c2e21b7b3c
8 changed files with 214 additions and 100 deletions

View file

@ -102,29 +102,58 @@ function matchesSource(entity: Record<string, unknown>, source: AggregateDefinit
* HNSWNounWithMetadata shape contract (standard fields top-level,
* custom fields in metadata) in a single source of truth.
*/
function computeGroupKey(
/**
* Compute the group key(s) an entity contributes to.
*
* Returns one key normally, but **multiple** when a groupBy dimension is an `unnest` field:
* the entity contributes once per distinct array element (cartesian product across multiple
* unnest dimensions). An entity whose unnest field is missing/empty contributes to no group
* (returns `[]`). This is the fan-out behind tag-frequency style aggregates.
*/
function computeGroupKeys(
entity: Record<string, unknown>,
groupBy: GroupByDimension[]
): Record<string, string | number> {
const key: Record<string, string | number> = {}
): Record<string, string | number>[] {
const e = entity as unknown as HNSWNounWithMetadata
let keys: Record<string, string | number>[] = [{}]
for (const dim of groupBy) {
if (typeof dim === 'string') {
const val = resolveEntityField(e, dim)
key[dim] = val !== undefined && val !== null ? String(val) : '__null__'
const v = val !== undefined && val !== null ? String(val) : '__null__'
for (const k of keys) k[dim] = v
} else if ('unnest' in dim) {
const val = resolveEntityField(e, dim.field)
const raw = Array.isArray(val) ? val : val !== undefined && val !== null ? [val] : []
// Distinct elements: an entity with duplicate tags counts once per distinct tag.
const elems = Array.from(new Set(raw.map(x => String(x))))
if (elems.length === 0) return [] // contributes to no group
const next: Record<string, string | number>[] = []
for (const k of keys) {
for (const el of elems) next.push({ ...k, [dim.field]: el })
}
keys = next
} else {
// Time-windowed field
const val = resolveEntityField(e, dim.field)
if (typeof val === 'number') {
key[dim.field] = bucketTimestamp(val, dim.window)
} else {
key[dim.field] = '__null__'
}
const v = typeof val === 'number' ? bucketTimestamp(val, dim.window) : '__null__'
for (const k of keys) k[dim.field] = v
}
}
return key
return keys
}
/**
* Single representative group key (first of {@link computeGroupKeys}). Retained for the
* materializer and back-compat; the incremental contribution paths use `computeGroupKeys`
* so unnest dimensions fan out correctly.
*/
function computeGroupKey(
entity: Record<string, unknown>,
groupBy: GroupByDimension[]
): Record<string, string | number> {
return computeGroupKeys(entity, groupBy)[0] ?? {}
}
/**
@ -451,40 +480,8 @@ export class AggregationIndex {
continue
}
const groupKey = computeGroupKey(entity, def.groupBy)
const serialized = serializeGroupKey(groupKey)
const stateMap = this.states.get(name)!
let group = stateMap.get(serialized)
if (!group) {
group = {
groupKey,
metrics: {},
lastUpdated: Date.now()
}
// Initialize all metrics
for (const metricName of Object.keys(def.metrics)) {
group.metrics[metricName] = freshMetricState()
}
stateMap.set(serialized, group)
}
// Increment each metric
for (const [metricName, metricDef] of Object.entries(def.metrics)) {
const state = group.metrics[metricName]
if (metricDef.op === 'count') {
state.count++
state.sum++
} else {
const val = getNumericField(entity, metricDef.field!)
if (val !== undefined) {
updateMetricAdd(state, val, metricDef.op)
}
}
}
group.lastUpdated = Date.now()
this.dirty.add(name)
// Shared with onEntityUpdated/backfill; fans out unnest dimensions to N groups.
this.addContribution(name, def, entity)
}
}
@ -661,37 +658,41 @@ export class AggregationIndex {
def: AggregateDefinition,
entity: Record<string, unknown>
): void {
const groupKey = computeGroupKey(entity, def.groupBy)
const serialized = serializeGroupKey(groupKey)
const stateMap = this.states.get(aggName)!
let group = stateMap.get(serialized)
if (!group) {
group = {
groupKey,
metrics: {},
lastUpdated: Date.now()
}
for (const metricName of Object.keys(def.metrics)) {
group.metrics[metricName] = freshMetricState()
}
stateMap.set(serialized, group)
}
// Fan out: an unnest dimension makes one entity contribute to several groups.
for (const groupKey of computeGroupKeys(entity, def.groupBy)) {
const serialized = serializeGroupKey(groupKey)
let group = stateMap.get(serialized)
for (const [metricName, metricDef] of Object.entries(def.metrics)) {
const state = group.metrics[metricName]
if (metricDef.op === 'count') {
state.count++
state.sum++
} else {
const val = getNumericField(entity, metricDef.field!)
if (val !== undefined) {
updateMetricAdd(state, val, metricDef.op)
if (!group) {
group = {
groupKey,
metrics: {},
lastUpdated: Date.now()
}
for (const metricName of Object.keys(def.metrics)) {
group.metrics[metricName] = freshMetricState()
}
stateMap.set(serialized, group)
}
for (const [metricName, metricDef] of Object.entries(def.metrics)) {
const state = group.metrics[metricName]
if (metricDef.op === 'count') {
state.count++
state.sum++
} else {
const val = getNumericField(entity, metricDef.field!)
if (val !== undefined) {
updateMetricAdd(state, val, metricDef.op)
}
}
}
group.lastUpdated = Date.now()
}
group.lastUpdated = Date.now()
this.dirty.add(aggName)
}
@ -704,40 +705,42 @@ export class AggregationIndex {
def: AggregateDefinition,
entity: Record<string, unknown>
): void {
const groupKey = computeGroupKey(entity, def.groupBy)
const serialized = serializeGroupKey(groupKey)
const stateMap = this.states.get(aggName)!
const group = stateMap.get(serialized)
if (!group) return
// Fan out: reverse the entity's contribution from every group it joined.
for (const groupKey of computeGroupKeys(entity, def.groupBy)) {
const serialized = serializeGroupKey(groupKey)
const group = stateMap.get(serialized)
if (!group) continue
for (const [metricName, metricDef] of Object.entries(def.metrics)) {
const state = group.metrics[metricName]
if (metricDef.op === 'count') {
state.count = Math.max(0, state.count - 1)
state.sum = Math.max(0, state.sum - 1)
} else {
const val = getNumericField(entity, metricDef.field!)
if (val !== undefined) {
updateMetricRemove(state, val, metricDef.op)
for (const [metricName, metricDef] of Object.entries(def.metrics)) {
const state = group.metrics[metricName]
if (metricDef.op === 'count') {
state.count = Math.max(0, state.count - 1)
state.sum = Math.max(0, state.sum - 1)
} else {
const val = getNumericField(entity, metricDef.field!)
if (val !== undefined) {
updateMetricRemove(state, val, metricDef.op)
// MIN/MAX can't be decremented — mark as potentially stale
if (val <= state.min || val >= state.max) {
if (!this.staleMinMax.has(aggName)) {
this.staleMinMax.set(aggName, new Set())
// MIN/MAX can't be decremented — mark as potentially stale
if (val <= state.min || val >= state.max) {
if (!this.staleMinMax.has(aggName)) {
this.staleMinMax.set(aggName, new Set())
}
this.staleMinMax.get(aggName)!.add(`${serialized}:${metricName}`)
}
this.staleMinMax.get(aggName)!.add(`${serialized}:${metricName}`)
}
}
}
}
// Remove group if all metrics are empty
const allEmpty = Object.values(group.metrics).every(m => m.count === 0)
if (allEmpty) {
stateMap.delete(serialized)
} else {
group.lastUpdated = Date.now()
// Remove group if all metrics are empty
const allEmpty = Object.values(group.metrics).every(m => m.count === 0)
if (allEmpty) {
stateMap.delete(serialized)
} else {
group.lastUpdated = Date.now()
}
}
this.dirty.add(aggName)

View file

@ -34,6 +34,7 @@
import type { Brainy } from '../brainy.js'
import type { NounType } from '../types/graphTypes.js'
import type { Vector } from '../coreTypes.js'
import { ExactMatchSignal } from './signals/ExactMatchSignal.js'
import { PatternSignal } from './signals/PatternSignal.js'
import { EmbeddingSignal } from './signals/EmbeddingSignal.js'
@ -211,6 +212,8 @@ export class SmartExtractor {
formatContext?: FormatContext
allTerms?: string[]
metadata?: any
/** Pre-computed candidate embedding (from a batch embed) — forwarded to EmbeddingSignal. */
vector?: Vector
},
minConfidence?: number
): Promise<ExtractionResult | null> {
@ -243,7 +246,8 @@ export class SmartExtractor {
const enrichedContext = {
definition: context?.definition,
allTerms: [...(context?.allTerms || []), ...formatHints],
metadata: context?.metadata
metadata: context?.metadata,
vector: context?.vector
}
// Execute all signals in parallel

View file

@ -137,7 +137,27 @@ export class NeuralEntityExtractor {
// Step 1: Extract potential entities using patterns
const candidates = await this.extractCandidates(text)
// Step 1b: Batch-embed the unique candidate texts once, instead of one brain.embed()
// per candidate inside EmbeddingSignal (N sequential model calls). The vectors are
// forwarded into the signal; if the batch fails, the signal falls back to per-candidate
// embedding, so this is a pure performance optimization with no behavior change.
const vectorByText = new Map<string, Vector>()
if (useNeuralMatching) {
const uniqueTexts = [...new Set(candidates.map(c => c.text))]
if (uniqueTexts.length > 0) {
try {
const vectors = await this.brain.embedBatch(uniqueTexts)
uniqueTexts.forEach((t, i) => {
const v = vectors[i]
if (v) vectorByText.set(t, v)
})
} catch {
// Leave the map empty — EmbeddingSignal will embed per-candidate as before.
}
}
}
// Step 2: Classify each candidate using SmartExtractor
for (const candidate of candidates) {
// Use SmartExtractor for unified neural + rule-based classification.
@ -146,7 +166,8 @@ export class NeuralEntityExtractor {
// option had no loosening effect).
const classification = await this.smartExtractor.extract(candidate.text, {
definition: candidate.context,
allTerms: [candidate.text, candidate.context]
allTerms: [candidate.text, candidate.context],
vector: vectorByText.get(candidate.text)
}, minConfidence)
// SmartExtractor already gates at minConfidence; this guards against null only.

View file

@ -151,6 +151,8 @@ export class EmbeddingSignal {
definition?: string
allTerms?: string[]
metadata?: any
/** Pre-computed candidate embedding (from a batch embed) — skips the per-candidate embed. */
vector?: Vector
}
): Promise<TypeSignal | null> {
this.stats.calls++
@ -167,8 +169,8 @@ export class EmbeddingSignal {
}
try {
// Embed candidate once (efficiency!)
const vector = await this.embedWithTimeout(candidate)
// Use the pre-computed vector when the caller batch-embedded; otherwise embed once here.
const vector = context?.vector ?? await this.embedWithTimeout(candidate)
// Check all three sources in parallel
const [typeMatch, graphMatch, historyMatch] = await Promise.all([

View file

@ -706,9 +706,17 @@ export type TimeWindowGranularity =
| { seconds: number }
/**
* A GROUP BY dimension either a plain metadata field or a time-windowed field
* A GROUP BY dimension one of:
* - a plain metadata field name (string),
* - a time-windowed field (`{ field, window }`), or
* - an **unnest** field (`{ field, unnest: true }`): the field holds an array, and the entity
* contributes once to a group per distinct element (e.g. `tags: string[]` tag frequency).
* An entity whose unnest field is missing or an empty array contributes to no group.
*/
export type GroupByDimension = string | { field: string; window: TimeWindowGranularity }
export type GroupByDimension =
| string
| { field: string; window: TimeWindowGranularity }
| { field: string; unnest: true }
/**
* Source filter for which entities feed into an aggregate