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

@ -11,6 +11,39 @@ Collective. The SDK wraps it — most products never call Brainy directly. Read
---
## v7.24.0 — 2026-05-26
**Affected products:** aggregation/reporting users + anyone calling `extractEntities` on
multi-entity text. Additive; drop-in from 7.23.x.
### New — array-unnest `groupBy` (tag frequency / faceted counts)
A `groupBy` dimension can now be `{ field, unnest: true }`: the field holds an array and the
entity contributes once per **distinct** element. Enables tag-frequency / label-count style
aggregates:
```typescript
brain.defineAggregate({
name: 'tag_frequency',
source: { type: NounType.Document },
groupBy: [{ field: 'tags', unnest: true }],
metrics: { count: { op: 'count' } }
})
// queryAggregate('tag_frequency', { orderBy: 'count', order: 'desc' })
```
Duplicate elements on one entity count once; an entity with an empty/missing array joins no group.
### Perf — entity extraction batch-embeds candidates
`extractEntities` / `extractConcepts` now embed all unique candidate spans in a single
`embedBatch` call instead of one `embed()` per candidate (N sequential model calls before). No
behavior change — and with vectors reliably available, the embedding signal more consistently
reinforces correct types (e.g. clearer Organization confidence). Falls back to per-candidate
embedding if a batch call fails.
---
## v7.23.0 — 2026-05-26
**Affected products:** anyone using aggregation (stats/dashboards), graph traversal, or entity

View file

@ -200,6 +200,25 @@ brain.defineAggregate({
// { region: 'EU', date: '2024-01' }
```
### Array Fields (Unnest)
Group by **each element** of an array-valued field — for tag frequencies, label counts, and
faceted breakdowns. Mark the dimension `{ field, unnest: true }`:
```typescript
brain.defineAggregate({
name: 'tag_frequency',
source: { type: NounType.Document },
groupBy: [{ field: 'tags', unnest: true }],
metrics: { count: { op: 'count' } }
})
// A document tagged ['ml', 'ai'] contributes once to the 'ml' group and once to 'ai'.
// Duplicate tags on one entity count once; an entity with no tags joins no group.
const top = await brain.queryAggregate('tag_frequency', { orderBy: 'count', order: 'desc' })
// [ { groupKey: { tags: 'ai' }, metrics: { count: 3 }, count: 3 }, ... ]
```
## Querying Aggregates
Aggregate results are queried through the standard `find()` method.

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

View file

@ -230,4 +230,28 @@ describe('BR-ADV-FEATURES-BUN regression', () => {
await b.close()
})
})
describe('R2 — array-unnest groupBy (tag frequency)', () => {
it('groups by each distinct array element; dedups per entity; skips empty', async () => {
const b: any = new Brainy({ storage: { type: 'memory' } })
await b.init()
b.defineAggregate({
name: 'tags',
source: { type: NounType.Concept },
groupBy: [{ field: 'tags', unnest: true }],
metrics: { count: { op: 'count' } }
})
await b.add({ data: 'a', type: NounType.Concept, metadata: { tags: ['ml', 'ai'] } })
await b.add({ data: 'b', type: NounType.Concept, metadata: { tags: ['ai', 'rust'] } })
await b.add({ data: 'c', type: NounType.Concept, metadata: { tags: ['ai'] } })
await b.add({ data: 'd', type: NounType.Concept, metadata: { tags: ['rust', 'rust'] } }) // dup → once
await b.add({ data: 'e', type: NounType.Concept, metadata: {} }) // no tags → no group
await b.flush()
const rows = await b.queryAggregate('tags', { orderBy: 'count', order: 'desc' })
const freq = Object.fromEntries(rows.map((r: any) => [r.groupKey.tags, r.count]))
expect(freq).toEqual({ ai: 3, rust: 2, ml: 1 })
await b.close()
})
})
})