Source File
	project.go

Belonging Package
	github.com/freeeve/gochickpeas/gql/internal/exec

// Non-aggregated projection: the terminal sink evaluates the output
// columns per pushed row, applies DISTINCT on arrival (before ORDER
// BY/LIMIT, as the standard requires -- first occurrence kept), then
// finalize sorts and paginates. Projected rows live in a chunked arena;
// the matched row is retained (arena-copied) only when an ORDER BY key is
// not a projected column. The Rust bounded top-k heap remains a possible
// follow-up with byte-identical results.
package exec

import (
	"maps"
	"slices"

	"github.com/freeeve/gochickpeas/flatset"
	"github.com/freeeve/gochickpeas/gql/internal/ast"
	"github.com/freeeve/gochickpeas/gql/internal/eval"
	"github.com/freeeve/gochickpeas/gql/internal/plan"
	"github.com/freeeve/gochickpeas/gql/internal/semantics"
	"github.com/freeeve/gochickpeas/gql/value"
)

// projSink is the non-aggregated terminal sink.
type projSink struct {
	ctx     *eval.Ctx
	proj    *plan.ProjPlan
	slots   map[string]int
	returns []RowEval
	outs    [][]value.Value
	oArena  rowArena
	// needM: an ORDER BY key evaluates over the matched row, so matched
	// rows must be retained alongside their projections.
	needM  bool
	ms     [][]value.Value
	mArena rowArena
	// DISTINCT state: a single output column dedups through distinctSet
	// (the u32 entity-id fast path for node/rel values, AppendKey bytes
	// otherwise); a multi-column row keys on the concatenated AppendKey
	// encoding -- both dedups thus share one canonical value encoding.
	seenOne *distinctSet
	seen    *flatset.ByteSet
	key     []byte
	// topk is the streaming bounded accumulator under ORDER BY + LIMIT:
	// at most skip+limit rows are retained, ordered by the sort's exact
	// total order (keys, then arrival), so a rejected row costs one
	// comparison and leaves nothing in the arenas. limitCap is the
	// LIMIT-without-ORDER-BY retention cap (arrival order IS the output
	// order, so rows past skip+limit can never surface).
	topk     *topKRows
	limitCap int
	// topk key evaluation state (mirrors sortRowsByOrder's scope).
	kColIdx []int
	kScope  map[string]int
	kBase   int
	kRowbuf []value.Value
	kBuf    []value.Value
	// kGate marks the payload-gated top-k path: every ORDER BY key either
	// is an output column or evaluates against the matched row alone
	// (kRowEval, the RETURN m ... ORDER BY m.x identity family), and no
	// DISTINCT dedups on the full row -- so the keys can evaluate first
	// and the heap can refuse a candidate BEFORE its remaining columns are
	// built (most candidates die on one comparison once the heap fills).
	kGate    bool
	kRowEval []RowEval
}

// topkPayloadBuilds counts full projected payload constructions on the
// gated top-k path -- the load-independent oracle: ORDER BY v ASC LIMIT k
// over ascending input builds exactly k payloads, not one per row.
// disableTopkGate pins differential tests to the unguarded
// build-then-offer flow (the reference the gate must match exactly).
var (
	topkPayloadBuilds int
	disableTopkGate   bool
)

func newProjSink(ctx *eval.Ctx, proj *plan.ProjPlan, slots map[string]int, width int) *projSink {
	p := &projSink{
		ctx: ctx, proj: proj, slots: slots,
		returns:  make([]RowEval, len(proj.Returns)),
		oArena:   rowArena{width: len(proj.Returns)},
		limitCap: -1,
	}
	for i, r := range proj.Returns {
		p.returns[i] = compileEval(ctx, r.Expr, slots)
	}
	for i := range proj.OrderBy {
		if plan.OrderColIndex(proj.OrderBy[i].Expr, proj.Columns, proj.Returns) < 0 {
			p.needM = true
			p.mArena = rowArena{width: width}
			break
		}
	}
	if proj.Distinct {
		if len(proj.Returns) == 1 {
			p.seenOne = &distinctSet{}
		} else {
			p.seen = &flatset.ByteSet{}
		}
	}
	if bound := orderBound(proj); bound >= 0 {
		// Retention is bounded, so size arena chunks to it (plus eviction
		// turnover) instead of the default sweep-sized chunk -- a
		// LIMIT-bounded sink otherwise pays a near-empty 16K-value chunk
		// per arena (the aggregate finalize precedent).
		p.oArena.chunkValues = min(arenaChunkValues, max(64, 2*bound)*len(proj.Returns))
		if p.needM {
			p.mArena.chunkValues = min(arenaChunkValues, max(64, 2*bound)*width)
		}
		if nk := len(proj.OrderBy); nk > 0 {
			p.topk = newTopKRows(bound, nk, proj.OrderBy)
			p.kColIdx = make([]int, nk)
			for k := range proj.OrderBy {
				p.kColIdx[k] = plan.OrderColIndex(proj.OrderBy[k].Expr, proj.Columns, proj.Returns)
			}
			p.kScope = make(map[string]int, len(slots)+len(proj.Columns))
			maps.Copy(p.kScope, slots)
			p.kBase = 0
			if p.needM {
				p.kBase = width
			}
			for i, c := range proj.Columns {
				p.kScope[c] = p.kBase + i
			}
			p.kBuf = make([]value.Value, nk)
			p.kGate = !proj.Distinct && !disableTopkGate
			for k := range proj.OrderBy {
				if p.kColIdx[k] < 0 {
					// Not an output column, but still gateable when the
					// key's references all resolve in the matched row's
					// scope (RETURN m ... ORDER BY m.x and friends): it
					// then evaluates pre-build, exactly like a column key.
					if semantics.CheckRefs(proj.OrderBy[k].Expr, slots) == nil {
						if p.kRowEval == nil {
							p.kRowEval = make([]RowEval, nk)
						}
						p.kRowEval[k] = compileEval(ctx, proj.OrderBy[k].Expr, slots)
						continue
					}
					p.kGate = false // key needs the built row
					break
				}
			}
		} else {
			// No ORDER BY: arrival order is output order, so retention
			// past skip+limit can never surface.
			p.limitCap = bound
		}
	}
	return p
}

// pushKeys evaluates the ORDER BY key vector for a just-projected row
// into p.kBuf (out is the projected row, row the matched row).
func (p *projSink) pushKeys(out, row []value.Value) []value.Value {
	built := false
	for k := range p.proj.OrderBy {
		if idx := p.kColIdx[k]; idx >= 0 {
			p.kBuf[k] = out[idx]
			continue
		}
		if !built {
			p.kRowbuf = append(p.kRowbuf[:0], row...)
			p.kRowbuf = append(p.kRowbuf, out...)
			built = true
		}
		p.kBuf[k] = eval.Eval(p.ctx, p.proj.OrderBy[k].Expr, p.kRowbuf, p.kScope)
	}
	return p.kBuf
}

func (p *projSink) push(row []value.Value) bool {
	if p.limitCap >= 0 && len(p.outs) >= p.limitCap {
		// Arrival order is output order (no ORDER BY), so nothing a
		// producer could still emit can surface: stop the walk.
		return false
	}
	// Gated top-k: evaluate only the key columns and ask the heap first;
	// a refused candidate never builds its remaining columns (nor touches
	// the arena). Skipping is byte-identical to offer-then-pop (see
	// wouldAccept), so the surviving rows match the unguarded path
	// exactly.
	if p.kGate {
		for k := range p.proj.OrderBy {
			if idx := p.kColIdx[k]; idx >= 0 {
				p.kBuf[k] = p.returns[idx].Eval(p.ctx, row, p.slots)
			} else {
				p.kBuf[k] = p.kRowEval[k].Eval(p.ctx, row, p.slots)
			}
		}
		if !p.topk.wouldAccept(p.kBuf) {
			p.topk.seq++ // rejected offers still order future arrivals
			return true
		}
		topkPayloadBuilds++
		out := p.oArena.alloc()
		for k := range p.proj.OrderBy {
			if idx := p.kColIdx[k]; idx >= 0 {
				out[idx] = p.kBuf[k]
			}
		}
		for i, c := range p.returns {
			if !p.kGateCol(i) {
				out[i] = c.Eval(p.ctx, row, p.slots)
			}
		}
		if !p.topk.offer(p.pushKeys(out, row), out) {
			p.oArena.rollback()
		}
		return true
	}
	out := p.oArena.alloc()
	for i, c := range p.returns {
		out[i] = c.Eval(p.ctx, row, p.slots)
	}
	if p.seenOne != nil {
		if !p.seenOne.add(out[0], &p.key) {
			p.oArena.rollback()
			return true
		}
	} else if p.seen != nil {
		p.key = p.key[:0]
		for _, v := range out {
			p.key = value.AppendKey(p.key, v)
		}
		if !p.seen.Add(p.key) {
			p.oArena.rollback()
			return true
		}
	}
	if p.topk != nil {
		// Keys evaluate here, against the live matched row -- so the
		// matched row never needs retaining at all on this path.
		if !p.topk.offer(p.pushKeys(out, row), out) {
			p.oArena.rollback()
		}
		return true
	}
	p.outs = append(p.outs, out)
	if p.needM {
		p.ms = append(p.ms, p.mArena.copyRow(row))
	}
	// Reaching the unordered cap exactly now also stops the walk (a
	// deferred false would cost one more full candidate).
	return p.limitCap < 0 || len(p.outs) < p.limitCap
}

// kGateCol reports whether output column i is one of the gated path's
// already-evaluated key columns.
func (p *projSink) kGateCol(i int) bool {
	for k := range p.kColIdx {
		if p.kColIdx[k] == i {
			return true
		}
	}
	return false
}

func (p *projSink) close() {}

// topKRows is the sink's bounded ORDER BY + LIMIT accumulator: at most
// bound rows survive, under the same total order finalize's sort applies
// (ORDER BY keys, then arrival sequence), so streaming selection equals
// materialize-sort-truncate exactly. A max-heap over parallel arrays;
// the worst survivor sits at the root and each rejected offer costs one
// comparison.
type topKRows struct {
	bound int
	nk    int
	desc  []bool
	keys  []value.Value // n*nk, admitted key vectors
	outs  [][]value.Value
	seqs  []int
	n     int
	seq   int
}

func newTopKRows(bound, nk int, order []ast.SortItem) *topKRows {
	t := &topKRows{bound: bound, nk: nk, desc: make([]bool, nk)}
	for i := range order {
		t.desc[i] = order[i].Desc
	}
	// The heap fills to bound and stays there; presizing (capped, so a
	// huge LIMIT over few offers does not overcommit) removes the append
	// regrowth of the three parallel arrays.
	if pre := min(bound, 1024); pre > 0 {
		t.keys = make([]value.Value, 0, pre*nk)
		t.outs = make([][]value.Value, 0, pre)
		t.seqs = make([]int, 0, pre)
	}
	return t
}

// cmpTo compares candidate (keys, seq) against admitted entry i.
func (t *topKRows) cmpTo(keys []value.Value, seq, i int) int {
	for k := 0; k < t.nk; k++ {
		ord := value.OrderCmp(keys[k], t.keys[i*t.nk+k])
		if t.desc[k] {
			ord = -ord
		}
		if ord != 0 {
			return ord
		}
	}
	return seq - t.seqs[i]
}

// cmpEntries compares admitted entries a and b.
func (t *topKRows) cmpEntries(a, b int) int {
	return t.cmpTo(t.keys[a*t.nk:(a+1)*t.nk], t.seqs[a], b)
}

// offer admits the row when it belongs in the current top bound, copying
// its keys (the caller's buffer is reused). Reports whether the row was
// retained.
// wouldAccept reports whether an offer carrying keys could be admitted,
// consulted BEFORE the payload is built: true under capacity; once full,
// only when keys sort strictly before the worst survivor. A key TIE
// loses -- a new candidate's arrival sequence is larger than every
// admitted row's, so a tied offer would be popped immediately -- which
// makes skipping a rejected candidate byte-identical to offering it, and
// preserves stability across a tie straddling the LIMIT boundary (the
// kept rows' relative order never changes).
func (t *topKRows) wouldAccept(keys []value.Value) bool {
	if t.bound == 0 {
		return false
	}
	if t.n < t.bound {
		return true
	}
	for k := 0; k < t.nk; k++ {
		ord := value.OrderCmp(keys[k], t.keys[k])
		if t.desc[k] {
			ord = -ord
		}
		if ord != 0 {
			return ord < 0
		}
	}
	return false
}

func (t *topKRows) offer(keys []value.Value, out []value.Value) bool {
	seq := t.seq
	t.seq++
	if t.bound == 0 {
		return false
	}
	if t.n == t.bound && t.cmpTo(keys, seq, 0) >= 0 {
		return false
	}
	if t.n < t.bound {
		t.keys = append(t.keys, keys...)
		t.outs = append(t.outs, out)
		t.seqs = append(t.seqs, seq)
		t.n++
		t.siftUp(t.n - 1)
		return true
	}
	copy(t.keys[:t.nk], keys)
	t.outs[0], t.seqs[0] = out, seq
	t.siftDown(0)
	return true
}

func (t *topKRows) swap(a, b int) {
	for k := 0; k < t.nk; k++ {
		t.keys[a*t.nk+k], t.keys[b*t.nk+k] = t.keys[b*t.nk+k], t.keys[a*t.nk+k]
	}
	t.outs[a], t.outs[b] = t.outs[b], t.outs[a]
	t.seqs[a], t.seqs[b] = t.seqs[b], t.seqs[a]
}

func (t *topKRows) siftUp(i int) {
	for i > 0 {
		parent := (i - 1) / 2
		if t.cmpEntries(i, parent) <= 0 {
			return
		}
		t.swap(i, parent)
		i = parent
	}
}

func (t *topKRows) siftDown(i int) {
	for {
		l := 2*i + 1
		if l >= t.n {
			return
		}
		big := l
		if r := l + 1; r < t.n && t.cmpEntries(r, l) > 0 {
			big = r
		}
		if t.cmpEntries(big, i) <= 0 {
			return
		}
		t.swap(i, big)
		i = big
	}
}

// sorted returns the survivors in final order (ascending under the total
// order).
func (t *topKRows) sorted() [][]value.Value {
	idx := make([]int, t.n)
	for i := range idx {
		idx[i] = i
	}
	slices.SortFunc(idx, t.cmpEntries)
	outs := make([][]value.Value, t.n)
	for i, j := range idx {
		outs[i] = t.outs[j]
	}
	return outs
}

func (p *projSink) finalize() [][]value.Value {
	if p.topk != nil {
		return paginate(p.topk.sorted(), p.proj.Skip, p.proj.Limit)
	}
	outs := p.outs
	if len(p.proj.OrderBy) > 0 {
		matchedAt := func(int) []value.Value { return nil }
		base := 0
		if p.needM {
			matchedAt = func(i int) []value.Value { return p.ms[i] }
			if len(p.ms) > 0 {
				base = len(p.ms[0])
			}
		}
		outs = sortRowsByOrder(p.ctx, p.proj, p.slots, matchedAt, base, outs)
	}
	return paginate(outs, p.proj.Skip, p.proj.Limit)
}

The pages are generated with Golds v0.8.7. (GOOS=linux GOARCH=amd64)
Golds is a Go 101 project developed by Tapir Liu.
PR and bug reports are welcome and can be submitted to the issue list.
Please follow @zigo_101 (reachable from the left QR code) to get the latest news of Golds.