fathom/internal/query/exec.go

package query

import (
	"fmt"
	"math"
	"sort"
	"strings"
	"sync"

	"git.z.gd/zzg/fathom/internal/model"
	"git.z.gd/zzg/fathom/internal/store"
)

// Result is a query's output.
type Result struct {
	Columns []string `json:"columns"`
	Rows    [][]any  `json:"rows"`

	// OrderBy names the measure the rows are sorted by. The UI charts this
	// column rather than guessing: a query that asks for the slowest tools and
	// gets a chart of call counts has answered a different question than the
	// one that was asked.
	OrderBy string     `json:"order_by,omitempty"`
	Stats   QueryStats `json:"stats"`
}

// QueryStats reports what the scan cost. Fathom has no index to explain a plan
// with, so the honest thing to show is how much it read.
type QueryStats struct {
	ShardsScanned int   `json:"shards_scanned"`
	ChunksScanned int   `json:"chunks_scanned"`
	RowsScanned   int   `json:"rows_scanned"`
	RowsMatched   int   `json:"rows_matched"`
	ElapsedMicros int64 `json:"elapsed_micros"`
}

// Engine executes queries against a Store.
type Engine struct {
	store *store.Store
}

func NewEngine(s *store.Store) *Engine { return &Engine{store: s} }

// Run executes a parsed query.
//
// Execution is one goroutine per shard fanning into a single merge — the
// collapsed form of Scuba's leaf/aggregator tree, minus the network. There is
// no plan to choose: the time range picks the shards, and everything else is a
// filter applied during the scan.
func (e *Engine) Run(q *Query) (*Result, error) {
	if q.Table != model.TableSpans && q.Table != model.TableLogs {
		return nil, fmt.Errorf("unknown table %q; Fathom has %q and %q", q.Table, model.TableSpans, model.TableLogs)
	}

	shards, err := e.store.ShardsOverlapping(q.Table, q.MinTS, q.MaxTS)
	if err != nil {
		return nil, err
	}

	var (
		mu      sync.Mutex
		stats   QueryStats
		acc     = newAccumulator(q)
		scanErr error
		wg      sync.WaitGroup
	)
	stats.ShardsScanned = len(shards)

	for _, info := range shards {
		wg.Add(1)
		go func(info store.ShardInfo) {
			defer wg.Done()
			part := newAccumulator(q)
			local, err := e.scanShard(q, info, part)

			mu.Lock()
			defer mu.Unlock()
			if err != nil && scanErr == nil {
				scanErr = err
			}
			stats.ChunksScanned += local.ChunksScanned
			stats.RowsScanned += local.RowsScanned
			stats.RowsMatched += local.RowsMatched
			acc.merge(part)
		}(info)
	}
	wg.Wait()

	if scanErr != nil {
		return nil, scanErr
	}
	res := acc.finish(q)
	res.Stats = stats
	res.OrderBy = q.OrderBy
	return res, nil
}

func (e *Engine) scanShard(q *Query, info store.ShardInfo, acc *accumulator) (QueryStats, error) {
	var stats QueryStats

	sh, err := e.store.OpenShardForRead(info.Table, info.Date)
	if err != nil {
		return stats, err
	}
	timeCol := store.TimeColumn(q.Table)

	err = sh.Chunks(func(c *store.Chunk) error {
		stats.ChunksScanned++
		stats.RowsScanned += c.Rows

		pred, skip := compile(q, c, timeCol)
		if skip {
			// An equality predicate whose literal is absent from this chunk's
			// dictionary cannot match any row in it. Skipping here is the only
			// pruning Fathom does below the shard level, and it is free.
			return nil
		}
		for i := 0; i < c.Rows; i++ {
			if !pred(i) {
				continue
			}
			stats.RowsMatched++
			acc.add(q, c, i)
		}
		return nil
	})
	return stats, err
}

// rowPred reports whether row i passes the WHERE clause.
type rowPred func(i int) bool

// compile builds a predicate closure bound to one chunk's physical layout.
//
// Binding per chunk is what keeps the inner loop cheap: string literals are
// resolved to dictionary codes once here, so the per-row work is integer
// comparison rather than string comparison. It returns skip=true when the chunk
// provably cannot match.
func compile(q *Query, c *store.Chunk, timeCol string) (pred rowPred, skip bool) {
	var preds []rowPred

	// The mandatory time range, applied per row. The shard granularity is a
	// day, so a range narrower than a day still has to filter within the shard.
	if tc, ok := c.Col(timeCol); ok && tc.Type == store.ColInt64 {
		min, max := q.MinTS, q.MaxTS
		vals := tc.I64
		preds = append(preds, func(i int) bool {
			return i < len(vals) && vals[i] >= min && vals[i] <= max
		})
	}

	for _, cond := range q.Conds {
		p, chunkSkip := compileCond(cond, c)
		if chunkSkip {
			return nil, true
		}
		if p != nil {
			preds = append(preds, p)
		}
	}

	switch len(preds) {
	case 0:
		return func(int) bool { return true }, false
	case 1:
		return preds[0], false
	}
	return func(i int) bool {
		for _, p := range preds {
			if !p(i) {
				return false
			}
		}
		return true
	}, false
}

func compileCond(cond Condition, c *store.Chunk) (rowPred, bool) {
	// Physical columns first.
	if col, ok := c.Col(cond.Column); ok {
		switch col.Type {
		case store.ColDict:
			return compileDictCond(cond, c, col)
		case store.ColInt64:
			vals := col.I64
			return func(i int) bool {
				return i < len(vals) && numCompare(float64(vals[i]), cond)
			}, false
		case store.ColUint8:
			return compileEnumCond(cond, col)
		case store.ColFloat32:
			vals := col.F32
			return func(i int) bool {
				return i < len(vals) && numCompare(float64(vals[i]), cond)
			}, false
		case store.ColFloat64:
			vals := col.F64
			return func(i int) bool {
				return i < len(vals) && numCompare(vals[i], cond)
			}, false
		}
	}
	// Otherwise it is an attribute, resolved through the side table.
	return compileAttrCond(cond, c), false
}

func compileDictCond(cond Condition, c *store.Chunk, col *store.Column) (rowPred, bool) {
	codes := col.Codes
	// Equality on a dictionary column is an integer compare once the literal is
	// resolved, and a whole-chunk skip if the literal is not in the dictionary.
	if cond.Op == OpEq {
		want, ok := c.Dict.Lookup(cond.Value)
		if !ok {
			return nil, true
		}
		return func(i int) bool { return i < len(codes) && codes[i] == want }, false
	}
	if cond.Op == OpNe {
		want, ok := c.Dict.Lookup(cond.Value)
		if !ok {
			return func(int) bool { return true }, false // nothing here equals it
		}
		return func(i int) bool { return i < len(codes) && codes[i] != want }, false
	}
	// Substring and pattern matching resolve the code per row. The dictionary
	// is small, so this stays cheap in practice.
	return func(i int) bool {
		if i >= len(codes) {
			return false
		}
		return strCompare(c.Dict.Get(codes[i]), cond)
	}, false
}

// compileEnumCond handles kind and status, which are stored as small integers
// but queried by name: `WHERE status = error`.
func compileEnumCond(cond Condition, col *store.Column) (rowPred, bool) {
	vals := col.U8
	var want uint8
	switch col.Name {
	case store.ColKind:
		want = uint8(model.ParseKind(cond.Value))
	case store.ColStatus:
		want = uint8(model.ParseStatus(cond.Value))
	default:
		if cond.IsNum {
			want = uint8(cond.Num)
		}
	}
	if cond.IsNum {
		return func(i int) bool { return i < len(vals) && numCompare(float64(vals[i]), cond) }, false
	}
	switch cond.Op {
	case OpEq:
		return func(i int) bool { return i < len(vals) && vals[i] == want }, false
	case OpNe:
		return func(i int) bool { return i < len(vals) && vals[i] != want }, false
	}
	return func(i int) bool { return i < len(vals) && numCompare(float64(vals[i]), cond) }, false
}

func compileAttrCond(cond Condition, c *store.Chunk) rowPred {
	keyCode, ok := c.Dict.Lookup(cond.Column)
	if !ok {
		// No row in this chunk has this attribute. `!=` is still satisfied by
		// absence, matching the intuition that a missing field is not equal to
		// the value you asked about.
		if cond.Op == OpNe {
			return func(int) bool { return true }
		}
		return func(int) bool { return false }
	}

	startCol, ok1 := c.Col("_attr_start")
	countCol, ok2 := c.Col("_attr_count")
	if !ok1 || !ok2 {
		return func(int) bool { return false }
	}
	starts, counts := startCol.I64, countCol.I64

	return func(i int) bool {
		if i >= len(starts) || i >= len(counts) {
			return false
		}
		start, n := int(starts[i]), int(counts[i])
		if start < 0 || start+n > len(c.AttrKeys) {
			return false
		}
		for j := start; j < start+n; j++ {
			if c.AttrKeys[j] != keyCode {
				continue
			}
			v := model.Value{
				Type: model.ValueType(c.AttrTypes[j]),
				Str:  c.Dict.Get(c.AttrVals[j]),
			}
			if cond.IsNum {
				if f, ok := v.Float(); ok {
					return numCompare(f, cond)
				}
			}
			return strCompare(v.Str, cond)
		}
		return cond.Op == OpNe
	}
}

func numCompare(v float64, c Condition) bool {
	if !c.IsNum {
		return false
	}
	switch c.Op {
	case OpEq:
		return v == c.Num
	case OpNe:
		return v != c.Num
	case OpLt:
		return v < c.Num
	case OpLe:
		return v <= c.Num
	case OpGt:
		return v > c.Num
	case OpGe:
		return v >= c.Num
	}
	return false
}

func strCompare(v string, c Condition) bool {
	switch c.Op {
	case OpEq:
		return v == c.Value
	case OpNe:
		return v != c.Value
	case OpContains:
		return strings.Contains(strings.ToLower(v), strings.ToLower(c.Value))
	case OpLike:
		return likeMatch(v, c.Value)
	case OpLt:
		return v < c.Value
	case OpLe:
		return v <= c.Value
	case OpGt:
		return v > c.Value
	case OpGe:
		return v >= c.Value
	}
	return false
}

// likeMatch implements SQL LIKE with % and _, case-insensitively.
func likeMatch(s, pattern string) bool {
	s, pattern = strings.ToLower(s), strings.ToLower(pattern)
	return likeHere(s, pattern)
}

func likeHere(s, p string) bool {
	for len(p) > 0 {
		switch p[0] {
		case '%':
			// Collapse runs of % so "a%%b" does not cost extra recursion.
			for len(p) > 0 && p[0] == '%' {
				p = p[1:]
			}
			if len(p) == 0 {
				return true
			}
			for i := 0; i <= len(s); i++ {
				if likeHere(s[i:], p) {
					return true
				}
			}
			return false
		case '_':
			if len(s) == 0 {
				return false
			}
			s, p = s[1:], p[1:]
		default:
			if len(s) == 0 || s[0] != p[0] {
				return false
			}
			s, p = s[1:], p[1:]
		}
	}
	return len(s) == 0
}

// value extracts a column or attribute from row i, for projection and grouping.
func value(c *store.Chunk, name string, i int) model.Value {
	if col, ok := c.Col(name); ok {
		switch col.Type {
		case store.ColDict:
			if i < len(col.Codes) {
				return model.StringValue(c.Dict.Get(col.Codes[i]))
			}
		case store.ColInt64:
			if i < len(col.I64) {
				return model.IntValue(col.I64[i])
			}
		case store.ColUint8:
			if i < len(col.U8) {
				switch name {
				case store.ColKind:
					return model.StringValue(model.Kind(col.U8[i]).String())
				case store.ColStatus:
					return model.StringValue(model.Status(col.U8[i]).String())
				}
				return model.IntValue(int64(col.U8[i]))
			}
		case store.ColFloat32:
			if i < len(col.F32) {
				return model.FloatValue(float64(col.F32[i]))
			}
		case store.ColFloat64:
			if i < len(col.F64) {
				return model.FloatValue(col.F64[i])
			}
		}
		return model.Value{}
	}
	// Attribute lookup.
	keyCode, ok := c.Dict.Lookup(name)
	if !ok {
		return model.Value{}
	}
	startCol, ok1 := c.Col("_attr_start")
	countCol, ok2 := c.Col("_attr_count")
	if !ok1 || !ok2 || i >= len(startCol.I64) {
		return model.Value{}
	}
	start, n := int(startCol.I64[i]), int(countCol.I64[i])
	if start < 0 || start+n > len(c.AttrKeys) {
		return model.Value{}
	}
	for j := start; j < start+n; j++ {
		if c.AttrKeys[j] == keyCode {
			return model.Value{
				Type: model.ValueType(c.AttrTypes[j]),
				Str:  c.Dict.Get(c.AttrVals[j]),
			}
		}
	}
	return model.Value{}
}

// sampleWeight is the multiplier a sampled row contributes to counts and sums.
//
// This is Scuba's sample-rate compensation: a row recorded at rate 0.01 stands
// for 100 rows, so counting it as 1 understates the total hundredfold. Averages
// and percentiles use the weight too, so a sampled dataset reports the same
// shape as an unsampled one.
func sampleWeight(c *store.Chunk, i int) float64 {
	col, ok := c.Col(store.ColSampleRate)
	if !ok || i >= len(col.F64) {
		return 1
	}
	r := col.F64[i]
	if r <= 0 || r > 1 {
		return 1
	}
	return 1 / r
}

// accumulator collects rows into either a row list or grouped aggregates.
type accumulator struct {
	grouped bool

	// Row-listing mode.
	rows [][]any

	// Aggregate mode, keyed by the joined group values.
	groups map[string]*group
	order  []string // first-seen order, so ungrouped output is stable
}

type group struct {
	key    []string
	aggs   []*aggState
	weight float64
}

type aggState struct {
	count   float64
	sum     float64
	min     float64
	max     float64
	seen    bool
	samples []float64 // retained only for percentiles
	weights []float64
}

func newAccumulator(q *Query) *accumulator {
	a := &accumulator{grouped: q.HasAggregates()}
	if a.grouped {
		a.groups = make(map[string]*group)
	}
	return a
}

func (a *accumulator) add(q *Query, c *store.Chunk, i int) {
	if !a.grouped {
		row := make([]any, len(q.Selects))
		for j, s := range q.Selects {
			if s.Column == "*" {
				row[j] = rowMap(c, i)
				continue
			}
			row[j] = jsonValue(value(c, s.Column, i))
		}
		a.rows = append(a.rows, row)
		return
	}

	key := make([]string, len(q.GroupBy))
	for j, g := range q.GroupBy {
		key[j] = value(c, g, i).Str
	}
	k := strings.Join(key, "\x00")

	g, ok := a.groups[k]
	if !ok {
		g = &group{key: key, aggs: make([]*aggState, len(q.Selects))}
		for j := range g.aggs {
			g.aggs[j] = &aggState{}
		}
		a.groups[k] = g
		a.order = append(a.order, k)
	}

	w := sampleWeight(c, i)
	g.weight += w

	for j, s := range q.Selects {
		if s.Agg == AggNone {
			continue
		}
		st := g.aggs[j]
		if s.Agg == AggCount && s.Column == "*" {
			st.count += w
			st.seen = true
			continue
		}
		v := value(c, s.Column, i)
		f, ok := v.Float()
		if !ok {
			// count over a non-numeric column still counts occurrences.
			if s.Agg == AggCount && v.Str != "" {
				st.count += w
				st.seen = true
			}
			continue
		}
		st.count += w
		st.sum += f * w
		if !st.seen || f < st.min {
			st.min = f
		}
		if !st.seen || f > st.max {
			st.max = f
		}
		st.seen = true
		if s.Agg.isPercentile() {
			st.samples = append(st.samples, f)
			st.weights = append(st.weights, w)
		}
	}
}

func (a *accumulator) merge(other *accumulator) {
	if !a.grouped {
		a.rows = append(a.rows, other.rows...)
		return
	}
	for _, k := range other.order {
		og := other.groups[k]
		g, ok := a.groups[k]
		if !ok {
			a.groups[k] = og
			a.order = append(a.order, k)
			continue
		}
		g.weight += og.weight
		for j := range og.aggs {
			if j >= len(g.aggs) {
				break
			}
			dst, src := g.aggs[j], og.aggs[j]
			if !src.seen {
				continue
			}
			if !dst.seen || src.min < dst.min {
				dst.min = src.min
			}
			if !dst.seen || src.max > dst.max {
				dst.max = src.max
			}
			dst.count += src.count
			dst.sum += src.sum
			dst.samples = append(dst.samples, src.samples...)
			dst.weights = append(dst.weights, src.weights...)
			dst.seen = true
		}
	}
}

func (a *accumulator) finish(q *Query) *Result {
	res := &Result{Columns: make([]string, len(q.Selects))}
	for i, s := range q.Selects {
		res.Columns[i] = s.Label()
	}

	if !a.grouped {
		res.Rows = a.rows
		sortRows(res, q)
		applyLimit(res, q)
		if res.Rows == nil {
			res.Rows = [][]any{}
		}
		return res
	}

	for _, k := range a.order {
		g := a.groups[k]
		row := make([]any, len(q.Selects))
		gi := 0
		for j, s := range q.Selects {
			if s.Agg == AggNone {
				if gi < len(g.key) {
					row[j] = g.key[gi]
				}
				gi++
				continue
			}
			row[j] = g.aggs[j].result(s.Agg)
		}
		res.Rows = append(res.Rows, row)
	}
	sortRows(res, q)
	applyLimit(res, q)
	if res.Rows == nil {
		res.Rows = [][]any{}
	}
	return res
}

func (s *aggState) result(f AggFunc) any {
	if !s.seen {
		if f == AggCount {
			return float64(0)
		}
		return nil
	}
	switch f {
	case AggCount:
		return s.count
	case AggSum:
		return s.sum
	case AggAvg:
		if s.count == 0 {
			return nil
		}
		return s.sum / s.count
	case AggMin:
		return s.min
	case AggMax:
		return s.max
	}
	if f.isPercentile() {
		return weightedPercentile(s.samples, s.weights, f.percentile())
	}
	return nil
}

// weightedPercentile computes a percentile honouring sample weights, so a
// sampled dataset yields the same distribution as an unsampled one.
func weightedPercentile(vals, weights []float64, p float64) any {
	if len(vals) == 0 {
		return nil
	}
	idx := make([]int, len(vals))
	for i := range idx {
		idx[i] = i
	}
	sort.Slice(idx, func(a, b int) bool { return vals[idx[a]] < vals[idx[b]] })

	var total float64
	for _, w := range weights {
		total += w
	}
	if total == 0 {
		total = float64(len(vals))
	}

	target := p * total
	var cum float64
	for _, i := range idx {
		w := 1.0
		if i < len(weights) {
			w = weights[i]
		}
		cum += w
		if cum >= target {
			return vals[i]
		}
	}
	return vals[idx[len(idx)-1]]
}

func sortRows(res *Result, q *Query) {
	if q.OrderBy == "" {
		// Default to newest first for row listings, which is what anyone
		// looking at a trace list actually wants.
		if !q.HasAggregates() {
			for i, c := range res.Columns {
				if c == store.ColStart || c == store.ColTS {
					sortByColumn(res, i, true)
					return
				}
			}
		}
		return
	}
	col := -1
	for i, c := range res.Columns {
		if strings.EqualFold(c, q.OrderBy) {
			col = i
			break
		}
	}
	if col < 0 {
		// ORDER BY a column that is not projected: fall back to leaving the
		// order alone rather than failing a query that is otherwise fine.
		return
	}
	sortByColumn(res, col, q.Desc)
}

func sortByColumn(res *Result, col int, desc bool) {
	sort.SliceStable(res.Rows, func(a, b int) bool {
		x, y := res.Rows[a][col], res.Rows[b][col]
		less := compareAny(x, y)
		if desc {
			return !less && compareAny(y, x)
		}
		return less
	})
}

func compareAny(a, b any) bool {
	af, aok := toFloat(a)
	bf, bok := toFloat(b)
	if aok && bok {
		return af < bf
	}
	return fmt.Sprint(a) < fmt.Sprint(b)
}

func toFloat(v any) (float64, bool) {
	switch t := v.(type) {
	case float64:
		if math.IsNaN(t) {
			return 0, false
		}
		return t, true
	case int64:
		return float64(t), true
	case int:
		return float64(t), true
	}
	return 0, false
}

func applyLimit(res *Result, q *Query) {
	if q.Limit > 0 && len(res.Rows) > q.Limit {
		res.Rows = res.Rows[:q.Limit]
	}
}

// jsonValue renders an attribute for JSON output with its original type, so a
// token count arrives in the UI as a number rather than a string.
func jsonValue(v model.Value) any {
	if v.Str == "" && v.Type == model.ValString {
		return ""
	}
	switch v.Type {
	case model.ValInt, model.ValFloat:
		if f, ok := v.Float(); ok {
			return f
		}
	case model.ValBool:
		return v.Str == "true"
	}
	return v.Str
}

// rowMap materialises a whole row for SELECT *.
func rowMap(c *store.Chunk, i int) map[string]any {
	out := map[string]any{}
	for j := range c.Cols {
		name := c.Cols[j].Name
		if strings.HasPrefix(name, "_") {
			continue // internal bookkeeping columns
		}
		out[name] = jsonValue(value(c, name, i))
	}
	return out
}