runtime: convert uniq operator to vio.Puller

compiler/rungen/op.go

@@ -18,7 +18,6 @@ import (
 	"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/runtime/sam/op/uniq"
 	"github.com/superdb/super/sbuf"
 	"github.com/superdb/super/vector"
 	"github.com/superdb/super/vector/vio"
@@ -109,8 +108,6 @@ func (b *Builder) compileLeaf(o dag.Op, parent sbuf.Puller) (sbuf.Puller, error)
 			return nil, err
 		}
 		return top.New(b.sctx(), parent, v.Limit, exprs, v.Reverse), nil
-	case *dag.UniqOp:
-		return uniq.New(b.rctx, parent, v.Cflag), nil
 	default:
 		return nil, fmt.Errorf("unknown DAG operator type: %v", v)
 	}

compiler/rungen/vop.go

@@ -411,6 +411,8 @@ 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.UniqOp:
+		return op.NewUniq(b.sctx(), parent, o.Cflag), nil
 	case *dag.UnnestOp:
 		e, err := b.compileVamExpr(o.Expr)
 		if err != nil {

runtime/op/uniq.go

@@ -0,0 +1,88 @@
+package op
+
+import (
+	"github.com/superdb/super"
+	"github.com/superdb/super/sbuf"
+	"github.com/superdb/super/scode"
+	"github.com/superdb/super/vector"
+	"github.com/superdb/super/vector/vio"
+)
+
+type Uniq struct {
+	sctx      *super.Context
+	parent    vio.Puller
+	withCount bool
+
+	count       int64
+	eos         bool
+	last        super.Value
+	pullBuilder scode.Builder
+	wrapBuilder scode.Builder
+}
+
+func NewUniq(sctx *super.Context, parent vio.Puller, withCount bool) *Uniq {
+	return &Uniq{
+		sctx:      sctx,
+		parent:    parent,
+		withCount: withCount,
+	}
+}
+
+func (u *Uniq) Pull(done bool) (vector.Any, error) {
+	if u.eos {
+		u.eos = false
+		return nil, nil
+	}
+	b := vector.NewDynamicValueBuilder()
+	for {
+		vec, err := u.parent.Pull(done)
+		if err != nil {
+			return nil, err
+		}
+		if vec == nil {
+			if u.count == 0 {
+				return nil, nil
+			}
+			val := u.wrap(u.last)
+			u.count = 0
+			u.eos = true
+			// The zero value is not a valid super.Value.
+			u.last = super.Value{}
+			return sbuf.Dematerialize(u.sctx, val), nil
+		}
+		for i := range vec.Len() {
+			u.writeUniq(b, vector.ValueAt(&u.pullBuilder, vec, i))
+		}
+		if vec := b.Build(u.sctx); vec.Len() > 0 {
+			return vec, nil
+		}
+	}
+}
+
+func (u *Uniq) writeUniq(b *vector.DynamicValueBuilder, val super.Value) {
+	if u.count == 0 {
+		u.count = 1
+		u.last = val.Copy()
+	} else if val.Equal(u.last) {
+		u.count++
+	} else {
+		b.Write(u.wrap(u.last))
+		u.count = 1
+		u.last = val.Copy()
+	}
+}
+
+func (u *Uniq) wrap(val super.Value) super.Value {
+	if !u.withCount {
+		return val
+	}
+	b := u.wrapBuilder
+	b.Truncate()
+	b.Append(val.Bytes())
+	b.Append(super.EncodeInt(u.count))
+	typ := u.sctx.MustLookupTypeRecord([]super.Field{
+		super.NewField("value", val.Type()),
+		super.NewField("count", super.TypeInt64),
+	})
+	return super.NewValue(typ, b.Bytes())
+}