unkin/pdbmux
4
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Combine aggregate columns per function instead of summing every one

Browse Source 
sumRows folded every numeric column by addition, which is only correct
for count and sum, so min/max returned a sum, avg an average of
averages, and a to_string extract collapsed into one empty-key row.

- Combine count and sum by adding, min and max by the extreme, on text
  columns as well as numeric ones
- Rewrite an avg extract into an upstream sum and count and divide the
  totals, answering under the avg key the client asked for
- Refuse an aggregate pdbmux cannot merge with 400 naming the clash
- Treat to_string and jsonb_typeof as row functions that group rather
  than fold, and key groups on every non-aggregate projected column
- Give the e2e fixture per-node resource line numbers and titles whose
  extremes differ per backend
This commit is contained in:
Unkin Agent
2026-09-06 23:17:10 +10:00
parent e889cf8f7f
commit 66ed7b615c
8 changed files with 1214 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aggregate.go
+318 -80
View File
@@ -2,83 +2,216 @@ package main
import (
"encoding/json"
"fmt"
"sort"
"strconv"
"strings"
)
// aggregateSpec names the columns of an extract/group_by result set: keys
// identify a row across backends, sums are the numeric aggregate columns added
// together.
type aggregateSpec struct {
keys []string
sums []string
// combineOp is how one aggregate column's per-backend values become one value.
type combineOp int
const (
opAdd combineOp = iota
opMin
opMax
)
// aggColumn is one aggregate column of a result row: the JSON key openvoxdb
// names it and the operation that folds the backends' values together.
type aggColumn struct {
name string
op combineOp
}
// parseAggregate reads a PuppetDB AST query and returns the aggregate shape of
// its response, or nil when the query is not an aggregate — only a top-level
// `extract` carrying at least one `["function", ...]` column produces summable
// rows. Key columns are the plain (non-function) extract fields, unioned with an
// explicit `group_by` clause when the query has one.
func parseAggregate(query string) *aggregateSpec {
// rowShape splits a result row into the columns that identify it and the
// aggregate columns that are folded. finish, when set, rewrites the merged row
// before it is encoded and forces re-encoding even for a single-backend row.
type rowShape struct {
keys []string
aggs []aggColumn
finish func(map[string]json.RawMessage)
}
// openvoxdb's accepted extract functions are the keys of pdb-fns->pg-fns,
// src/puppetlabs/puppetdb/query_eng/engine.clj:234-241, and each response column
// is named after the function itself (compile-fnexpression, engine.clj:1488-1495).
// count and sum add across backends; min and max take the extreme, on text
// columns as well as numeric ones, which openvoxdb allows because the numeric
// guard applies only to comparison clauses (engine.clj:2483-2490). avg is not
// combinable from the shard rows at all, so it is rewritten upstream — see
// rewriteAvg.
var aggregateOps = map[string]combineOp{
"count": opAdd,
"sum": opAdd,
"min": opMin,
"max": opMax,
}
// to_string and jsonb_typeof compile to scalar Postgres expressions, so they
// yield one row per input row rather than an aggregate: they name an ordinary
// projected column and identify a row instead of being folded into one.
var rowFns = map[string]bool{
"to_string": true,
"jsonb_typeof": true,
}
const avgColumn = "avg"
// aggregateSpec is the merge shape of an extract query's response.
type aggregateSpec struct {
keys []string
aggs []aggColumn
query string // rewritten upstream query, empty when the request's own is used
avg bool // the client asked for avg; sum and count are pdbmux's helpers
}
// parseAggregate reads a PuppetDB AST query and returns the merge shape of its
// response, or nil when the query is not an extract carrying a `["function",...]`
// column. A non-nil error means the query is an aggregate pdbmux cannot merge
// and must be refused rather than answered with a wrong number.
//
// Key columns are the plain extract fields, the row-function columns and any
// `group_by` clause; only genuine aggregates are folded.
func parseAggregate(query string) (*aggregateSpec, error) {
if strings.TrimSpace(query) == "" {
return nil
return nil, nil
}
var ast []json.RawMessage
if json.Unmarshal([]byte(query), &ast) != nil || len(ast) < 2 {
return nil
return nil, nil
}
var op string
if json.Unmarshal(ast[0], &op) != nil || op != "extract" {
return nil
return nil, nil
}
var cols []json.RawMessage
if json.Unmarshal(ast[1], &cols) != nil {
return nil
return nil, nil
}
spec := &aggregateSpec{}
var avgArgs []json.RawMessage
sawFunction := false
for _, col := range cols {
var name string
if json.Unmarshal(col, &name) == nil {
spec.keys = appendUnique(spec.keys, name)
continue
}
if fn, ok := functionName(col); ok {
spec.sums = appendUnique(spec.sums, fn)
fn, args, ok := functionColumn(col)
if !ok {
continue
}
sawFunction = true
switch {
case rowFns[fn]:
spec.keys = appendUnique(spec.keys, fn)
case fn == avgColumn:
if spec.avg {
return nil, fmt.Errorf("extract projects %q more than once, which openvoxdb answers with order-dependent columns", avgColumn)
}
if len(args) == 0 {
return nil, fmt.Errorf("extract function avg needs a column to average")
}
spec.avg, avgArgs = true, args
default:
combine, known := aggregateOps[fn]
if !known {
return nil, fmt.Errorf("extract function %q cannot be merged across backends", fn)
}
spec.aggs = appendUniqueAgg(spec.aggs, aggColumn{name: fn, op: combine})
}
}
if len(spec.sums) == 0 {
return nil
if !sawFunction {
return nil, nil
}
for _, node := range ast[2:] {
for _, f := range groupByFields(node) {
spec.keys = appendUnique(spec.keys, f)
}
}
return spec
if spec.avg {
if err := spec.rewriteAvg(ast, cols, avgArgs); err != nil {
return nil, err
}
}
return spec, nil
}
// functionName returns the response column an extract function produces, which
// PuppetDB names after the function itself: ["function","count","certname"]
// yields a "count" column.
func functionName(col json.RawMessage) (string, bool) {
// rewriteAvg replaces the client's avg column with the sum and count of the same
// expression, so the true weighted average can be computed from the shards:
// Postgres avg(x) is sum(x)/count(x), and both of those do combine.
func (a *aggregateSpec) rewriteAvg(ast, cols []json.RawMessage, args []json.RawMessage) error {
for _, taken := range []string{"sum", "count"} {
// openvoxdb names a column after its function, so a second one of the same
// name comes back as an order-dependent "<name>_2" rather than its own key.
if contains(a.keys, taken) || hasAgg(a.aggs, taken) {
return fmt.Errorf("avg cannot be merged alongside a %q column: pdbmux rewrites avg into an upstream sum and count, which openvoxdb would return under the same response key", taken)
}
}
sumCol, err := functionNode("sum", args)
if err != nil {
return err
}
countCol, err := functionNode("count", args)
if err != nil {
return err
}
out := make([]json.RawMessage, 0, len(cols)+1)
for _, col := range cols {
if fn, _, ok := functionColumn(col); ok && fn == avgColumn {
out = append(out, sumCol)
continue
}
out = append(out, col)
}
out = append(out, countCol)
rewritten := append([]json.RawMessage(nil), ast...)
encoded, err := json.Marshal(out)
if err != nil {
return err
}
rewritten[1] = encoded
query, err := json.Marshal(rewritten)
if err != nil {
return err
}
a.query = string(query)
a.aggs = append(a.aggs, aggColumn{name: "sum", op: opAdd}, aggColumn{name: "count", op: opAdd})
return nil
}
func functionNode(fn string, args []json.RawMessage) (json.RawMessage, error) {
parts := make([]json.RawMessage, 0, len(args)+2)
parts = append(parts, json.RawMessage(`"function"`), json.RawMessage(`"`+fn+`"`))
parts = append(parts, args...)
return json.Marshal(parts)
}
// functionColumn returns the response column an extract function produces, which
// openvoxdb names after the function itself — ["function","count","certname"]
// yields a "count" column — along with the function's arguments.
func functionColumn(col json.RawMessage) (string, []json.RawMessage, bool) {
var parts []json.RawMessage
if json.Unmarshal(col, &parts) != nil || len(parts) < 2 {
return "", false
return "", nil, false
}
var head, name string
if json.Unmarshal(parts[0], &head) != nil || head != "function" {
return "", false
return "", nil, false
}
if json.Unmarshal(parts[1], &name) != nil || name == "" {
return "", false
return "", nil, false
}
return name, true
return name, parts[2:], true
}
// groupByFields returns the field names of a ["group_by", ...] AST node, or nil
// for any other node.
// groupByFields returns the response columns a ["group_by", ...] AST node names,
// or nil for any other node. An entry may be a plain field or a function node,
// which groups on that function's own column.
func groupByFields(node json.RawMessage) []string {
var parts []json.RawMessage
if json.Unmarshal(node, &parts) != nil || len(parts) < 2 {
@@ -93,6 +226,10 @@ func groupByFields(node json.RawMessage) []string {
var name string
if json.Unmarshal(p, &name) == nil {
out = append(out, name)
continue
}
if fn, _, ok := functionColumn(p); ok {
out = append(out, fn)
}
}
return out
@@ -105,19 +242,53 @@ func appendUnique(s []string, v string) []string {
return append(s, v)
}
// columns reports which fields of a row form its grouping key and which are
// summed. A spec is fixed by the query, so the row is ignored.
func (a *aggregateSpec) columns(map[string]json.RawMessage) ([]string, []string) {
return a.keys, a.sums
func appendUniqueAgg(s []aggColumn, v aggColumn) []aggColumn {
if hasAgg(s, v.name) {
return s
}
return append(s, v)
}
// inferredColumns derives an event-counts row's shape from the row itself: the
func hasAgg(s []aggColumn, name string) bool {
for _, x := range s {
if x.name == name {
return true
}
}
return false
}
// shape reports a row's merge shape. A spec is fixed by the query, so the row is
// ignored.
func (a *aggregateSpec) shape(map[string]json.RawMessage) rowShape {
sh := rowShape{keys: a.keys, aggs: a.aggs}
if a.avg {
sh.finish = finishAvg
}
return sh
}
// finishAvg turns the summed helper columns into the avg key the client asked
// for. An empty count is openvoxdb's own answer for an average over no rows.
func finishAvg(row map[string]json.RawMessage) {
sum, sumOK := numberOf(row["sum"])
count, countOK := numberOf(row["count"])
delete(row, "sum")
delete(row, "count")
if !sumOK || !countOK || count == 0 {
row[avgColumn] = json.RawMessage("null")
return
}
row[avgColumn] = json.RawMessage(strconv.FormatFloat(sum/count, 'f', -1, 64))
}
// inferredShape derives an event-counts row's shape from the row itself: the
// counts to add (successes, failures, noops, skips, total) are its numeric
// fields, plus any null one — PuppetDB nulls an aggregate column when a backend
// fields, plus any null one — openvoxdb nulls an aggregate column when a backend
// matched nothing — and everything else, subject_type/subject/summarize_by,
// identifies the row. Those endpoints have a fixed response shape with no
// numeric key field, so nothing summable is mistaken for identity.
func inferredColumns(row map[string]json.RawMessage) ([]string, []string) {
func inferredShape(row map[string]json.RawMessage) rowShape {
var keys, sums []string
for name, val := range row {
if isJSONNumber(val) || isJSONNull(val) {
@@ -128,7 +299,11 @@ func inferredColumns(row map[string]json.RawMessage) ([]string, []string) {
}
sort.Strings(keys)
sort.Strings(sums)
return keys, sums
aggs := make([]aggColumn, 0, len(sums))
for _, s := range sums {
aggs = append(aggs, aggColumn{name: s, op: opAdd})
}
return rowShape{keys: keys, aggs: aggs}
}
// isJSONNumber reports whether a raw JSON value is a number.
@@ -144,31 +319,33 @@ func isJSONNull(raw json.RawMessage) bool {
return strings.TrimSpace(string(raw)) == "null"
}
// sumGroup accumulates the rows sharing one grouping key.
type sumGroup struct {
raw json.RawMessage // first contributing row, verbatim
row map[string]json.RawMessage // its decoded fields
totals map[string]float64 // running sum per aggregate column
merged bool // a second row was folded in
// mergeGroup accumulates the rows sharing one grouping key.
type mergeGroup struct {
raw json.RawMessage // first contributing row, verbatim
row map[string]json.RawMessage // its decoded fields
totals map[string]float64 // running sum per additive column
extremes map[string]json.RawMessage // running min/max per extreme column
finish func(map[string]json.RawMessage)
merged bool // a later row changed the group
}
// sumRows folds each backend's aggregate rows into one row per grouping key,
// adding the numeric aggregate columns. columns decides, per row, which fields
// are the key and which are summed.
// combineRows folds each backend's aggregate rows into one row per grouping key,
// combining each aggregate column by its own operation. shape decides, per row,
// which fields are the key and how each aggregate column combines.
//
// A row that is not a JSON object passes through untouched, as does the sole row
// of a key only one backend reported — those keep their upstream bytes. An
// aggregate column that is absent or non-numeric in a later row is left at the
// earlier backend's value rather than being coerced to zero. results come in
// configured backend order and the output keeps first-seen order, a tie-break
// only.
func sumRows(results []backendResult, columns func(map[string]json.RawMessage) ([]string, []string)) []json.RawMessage {
// A row that is not a JSON object passes through untouched, as does the sole
// row of a key only one backend reported and every row of a shape with no
// aggregate column to fold — those keep their upstream bytes. An aggregate
// column that is absent or null in a later row leaves the group's value alone
// rather than being coerced to zero. results come in configured backend order
// and the output keeps first-seen order, a tie-break only.
func combineRows(results []backendResult, shape func(map[string]json.RawMessage) rowShape) []json.RawMessage {
type slot struct {
raw json.RawMessage // passthrough row, when group is nil
group *sumGroup
group *mergeGroup
}
var order []slot
groups := map[string]*sumGroup{}
groups := map[string]*mergeGroup{}
for _, res := range results {
for _, rec := range res.records {
@@ -177,32 +354,22 @@ func sumRows(results []backendResult, columns func(map[string]json.RawMessage) (
order = append(order, slot{raw: rec.Raw})
continue
}
keys, sums := columns(row)
k := groupKey(row, keys)
sh := shape(row)
if len(sh.aggs) == 0 {
// Only row functions were projected, so there is nothing to fold
// and each backend's rows stand on their own.
order = append(order, slot{raw: rec.Raw})
continue
}
k := groupKey(row, sh.keys)
g, ok := groups[k]
if !ok {
g = &sumGroup{raw: rec.Raw, row: row, totals: map[string]float64{}}
for _, s := range sums {
if n, ok := numberOf(row[s]); ok {
g.totals[s] = n
}
}
g = newMergeGroup(rec.Raw, row, sh)
groups[k] = g
order = append(order, slot{group: g})
continue
}
for _, s := range sums {
n, ok := numberOf(row[s])
if !ok {
continue
}
if _, seen := g.totals[s]; !seen {
// First numeric value for a column the earlier row lacked.
g.totals[s] = 0
}
g.totals[s] += n
g.merged = true
}
g.fold(row, sh.aggs)
}
}
@@ -217,10 +384,75 @@ func sumRows(results []backendResult, columns func(map[string]json.RawMessage) (
return out
}
func newMergeGroup(raw json.RawMessage, row map[string]json.RawMessage, sh rowShape) *mergeGroup {
g := &mergeGroup{
raw: raw,
row: row,
totals: map[string]float64{},
extremes: map[string]json.RawMessage{},
finish: sh.finish,
}
for _, c := range sh.aggs {
if c.op == opAdd {
if n, ok := numberOf(row[c.name]); ok {
g.totals[c.name] = n
}
continue
}
if v, ok := row[c.name]; ok && !isJSONNull(v) {
g.extremes[c.name] = v
}
}
return g
}
func (g *mergeGroup) fold(row map[string]json.RawMessage, aggs []aggColumn) {
for _, c := range aggs {
if c.op == opAdd {
n, ok := numberOf(row[c.name])
if !ok {
continue
}
if _, seen := g.totals[c.name]; !seen {
// First numeric value for a column the earlier row lacked.
g.totals[c.name] = 0
}
g.totals[c.name] += n
g.merged = true
continue
}
v, ok := row[c.name]
if !ok || isJSONNull(v) {
continue
}
cur, seen := g.extremes[c.name]
if seen && !extremeWins(v, cur, c.op) {
continue
}
g.extremes[c.name] = v
g.merged = true
}
}
// extremeWins reports whether candidate replaces the running min or max. Values
// are compared decoded, so min/max works on the text columns openvoxdb allows
// them on as well as on numbers.
func extremeWins(candidate, current json.RawMessage, op combineOp) bool {
var a, b any
if json.Unmarshal(candidate, &a) != nil || json.Unmarshal(current, &b) != nil {
return false
}
c := compareValues(a, b)
if op == opMin {
return c < 0
}
return c > 0
}
// encode renders a group back to JSON, reusing the first row's bytes when
// nothing was added to it.
func (g *sumGroup) encode() json.RawMessage {
if !g.merged {
// nothing changed it and no finish step has to rewrite it.
func (g *mergeGroup) encode() json.RawMessage {
if !g.merged && g.finish == nil {
return g.raw
}
row := make(map[string]json.RawMessage, len(g.row))
@@ -230,6 +462,12 @@ func (g *sumGroup) encode() json.RawMessage {
for col, total := range g.totals {
row[col] = json.RawMessage(strconv.FormatFloat(total, 'f', -1, 64))
}
for col, v := range g.extremes {
row[col] = v
}
if g.finish != nil {
g.finish(row)
}
raw, err := json.Marshal(row)
if err != nil {
return g.raw