diff --git a/compiler/rungen/op.go b/compiler/rungen/op.go index 89b9e12fc8..5deffd712e 100644 --- a/compiler/rungen/op.go +++ b/compiler/rungen/op.go @@ -17,8 +17,6 @@ import ( "github.com/superdb/super/runtime/expr" "github.com/superdb/super/runtime/op" samexpr "github.com/superdb/super/runtime/sam/expr" - "github.com/superdb/super/runtime/sam/op/top" - "github.com/superdb/super/sbuf" "github.com/superdb/super/vector" "github.com/superdb/super/vector/vio" ) @@ -100,19 +98,6 @@ func (b *Builder) lookupType(id int) (super.Type, error) { return typ, nil } -func (b *Builder) compileLeaf(o dag.Op, parent sbuf.Puller) (sbuf.Puller, error) { - switch v := o.(type) { - case *dag.TopOp: - exprs, err := b.compileSortExprs(v.Exprs) - if err != nil { - return nil, err - } - return top.New(b.sctx(), parent, v.Limit, exprs, v.Reverse), nil - default: - return nil, fmt.Errorf("unknown DAG operator type: %v", v) - } -} - func (b *Builder) compileAssignmentsToLvals(assignments []dag.Assignment) ([]*samexpr.Lval, []*samexpr.Lval, error) { var srcs, dsts []*samexpr.Lval for _, a := range assignments { diff --git a/compiler/rungen/vop.go b/compiler/rungen/vop.go index 4024462537..473274d634 100644 --- a/compiler/rungen/vop.go +++ b/compiler/rungen/vop.go @@ -411,6 +411,12 @@ func (b *Builder) compileVamLeaf(o dag.Op, parent vio.Puller) (vio.Puller, error return op.NewTableScanner(b.rctx, slicer, table, nil, nil, b.progress), nil case *dag.TableMetaScan: return op.NewTableMetaScanner(b.rctx.Context, b.sctx(), b.env.DB(), o.ID, o.Meta) + case *dag.TopOp: + exprs, err := b.compileSortExprs(o.Exprs) + if err != nil { + return nil, err + } + return op.NewTop(b.sctx(), parent, o.Limit, exprs, o.Reverse), nil case *dag.UniqOp: return op.NewUniq(b.sctx(), parent, o.Cflag), nil case *dag.UnnestOp: @@ -426,15 +432,7 @@ func (b *Builder) compileVamLeaf(o dag.Op, parent vio.Puller) (vio.Puller, error } return op.NewValues(b.sctx(), parent, exprs), nil default: - var sbufParent sbuf.Puller - if parent != nil { - sbufParent = sbuf.NewMaterializer(parent) - } - sbufPuller, err := b.compileLeaf(o, sbufParent) - if err != nil { - return nil, err - } - return sbuf.NewDematerializer(b.sctx(), sbufPuller), nil + return nil, fmt.Errorf("unknown DAG operator %T", o) } } diff --git a/runtime/sam/op/top/top.go b/runtime/op/top.go similarity index 53% rename from runtime/sam/op/top/top.go rename to runtime/op/top.go index 9e8bfa1444..8401eed8e2 100644 --- a/runtime/sam/op/top/top.go +++ b/runtime/op/top.go @@ -1,30 +1,32 @@ -package top +package op import ( "container/heap" + "slices" "github.com/superdb/super" "github.com/superdb/super/runtime/sam/expr" "github.com/superdb/super/runtime/sam/op/sort" - "github.com/superdb/super/sbuf" + "github.com/superdb/super/scode" + "github.com/superdb/super/vector" + "github.com/superdb/super/vector/vio" ) -// Top produces the first N values that sort would produce with the same arguments. -type Op struct { +type Top struct { sctx *super.Context - parent sbuf.Puller + parent vio.Puller limit int exprs []expr.SortExpr guessReverse bool + builder scode.Builder + compare expr.CompareFn eos bool records *recordSlice - compare expr.CompareFn } -// New returns an operator that produces the first limit -func New(sctx *super.Context, parent sbuf.Puller, limit int, exprs []expr.SortExpr, guessReverse bool) *Op { - return &Op{ +func NewTop(sctx *super.Context, parent vio.Puller, limit int, exprs []expr.SortExpr, guessReverse bool) *Top { + return &Top{ sctx: sctx, parent: parent, limit: limit, @@ -33,57 +35,60 @@ func New(sctx *super.Context, parent sbuf.Puller, limit int, exprs []expr.SortEx } } -func (o *Op) Pull(done bool) (sbuf.Batch, error) { - if o.eos { - o.eos = false +func (t *Top) Pull(done bool) (vector.Any, error) { + if t.eos { + t.eos = false return nil, nil } for { - batch, err := o.parent.Pull(done) + vec, err := t.parent.Pull(done) if err != nil { return nil, err } - if batch == nil { - if o.records == nil { + if vec == nil { + if t.records == nil { return nil, nil } - o.eos = true - return o.sorted(), nil + t.eos = true + return t.sorted(), nil } - vals := batch.Values() - for i := range vals { - o.consume(vals[i]) + for i := range vec.Len() { + t.consume(vector.ValueAt(&t.builder, vec, i)) } - batch.Unref() } } -func (o *Op) consume(rec super.Value) { - if o.records == nil { - if o.compare == nil { - comparator := sort.NewComparator(o.sctx, o.exprs, rec, o.guessReverse) +func (t *Top) consume(val super.Value) { + if t.records == nil { + if t.compare == nil { + comparator := sort.NewComparator(t.sctx, t.exprs, val, t.guessReverse) // package heap implements a min-heap. Invert the comparison result to get a max-heap. - o.compare = func(a, b super.Value) int { return comparator.Compare(a, b) * -1 } + t.compare = func(a, b super.Value) int { return comparator.Compare(a, b) * -1 } } - o.records = newRecordSlice(o.compare) - heap.Init(o.records) + t.records = newRecordSlice(t.compare) + heap.Init(t.records) } - if o.records.Len() < o.limit || o.compare(o.records.Index(0), rec) < 0 { - heap.Push(o.records, rec.Copy()) + if t.records.Len() < t.limit || t.compare(t.records.Index(0), val) < 0 { + heap.Push(t.records, val.Copy()) } - if o.records.Len() > o.limit { - heap.Pop(o.records) + if t.records.Len() > t.limit { + heap.Pop(t.records) } } -func (o *Op) sorted() sbuf.Batch { - out := make([]super.Value, o.records.Len()) - for i := o.records.Len() - 1; i >= 0; i-- { - out[i] = heap.Pop(o.records).(super.Value) +func (t *Top) sorted() vector.Any { + // Transform t.record.vals from a max-heap to an ascending list that we + // can use to write values to a vector builder in ascending order. + vals := t.records.vals + for i := range slices.Backward(vals) { + vals[i] = heap.Pop(t.records).(super.Value) + } + t.records = nil + b := vector.NewDynamicValueBuilder() + for _, val := range vals { + b.Write(val) } - // clear records - o.records = nil - return sbuf.NewArray(out) + return b.Build(t.sctx) } type recordSlice struct {